Added more tests for decorator

tests/test_decorator.py

@@ -4,6 +4,27 @@
 from tombstones.decorator import line_number_for_tombstone
 
 
+@mock.patch('tombstones.decorator.save_log_entry')
+@mock.patch('tombstones.decorator.LogEntry')
+@mock.patch('datetime.datetime')
+def test_log_entry_arguments(mock_datetime, log_entry, mock_save_log_entry):
+    now = str(mock.MagicMock())
+    mock_datetime.now.return_value = now
+
+    @tombstone
+    def test_function():
+        pass
+    test_function()
+
+    assert log_entry.called
+    log_entry.assert_called_with(
+        name='test_function',
+        line_number=14,
+        datetime=now,
+        source_file=__file__,
+    )
+
+
 def test_tombstone_decorator():
     def function(test_value):
         return test_value
@@ -13,7 +34,7 @@ def function(test_value):
 
 
 @mock.patch('tombstones.decorator.save_log_entry')
-def test_save_log_entry(mock_save_log_entry):
+def test_decorated_function(mock_save_log_entry):
     @tombstone
     def function():
         pass
@@ -22,24 +43,22 @@ def function():
 
 
 @mock.patch('tombstones.decorator.save_log_entry')
-@mock.patch('tombstones.decorator.LogEntry')
-@mock.patch('datetime.datetime')
-def test_log_entry_arguments(mock_datetime, log_entry, mock_save_log_entry):
-    now = str(mock.MagicMock())
-    mock_datetime.now.return_value = now
-
+def test_decorated_class(mock_save_log_entry):
     @tombstone
-    def test_function():
+    class Example(object):
         pass
-    test_function()
+    Example()
+    assert mock_save_log_entry.called
 
-    assert log_entry.called
-    log_entry.assert_called_with(
-        name='test_function',
-        line_number=31,
-        datetime=now,
-        source_file=__file__,
-    )
+
+@mock.patch('tombstones.decorator.save_log_entry')
+def test_decorated_method(mock_save_log_entry):
+    class Example(object):
+        @tombstone
+        def example_method(self):
+            pass
+    Example().example_method()
+    assert mock_save_log_entry.called
 
 
 def test_first_line_number_for_tombstone():