Add magic method tracking

src/main/python/karellen/testing/mock/__init__.py

@@ -18,9 +18,9 @@
 
 import types
 
-from unittest.mock import MagicMock, DEFAULT
+from unittest.mock import patch, Mock, NonCallableMock, DEFAULT, _all_magics, _calculate_return_value
 
-__all__ = ["Spy", "MagicSpy"]
+__all__ = ["Spy", "MagicSpy", "instrument_wrapped_magic"]
 
 
 class Spy(object):
@@ -38,38 +38,59 @@ class Spy(object):
         calls.
 
         Obviously, Spy, being "self" in every method call, must provide for field access (get and set), which are also
-        passed through the
+        passed through this spy and are delegated as appropriate.
     """
 
+    def __new__(cls, *args, **kw):
+        # every instance has its own class
+        # so we can create magic methods on the
+        # class without stomping on other mocks
+        new = type(cls.__name__, (cls,), {'__doc__': cls.__doc__})
+        instance = object.__new__(new)
+        return instance
+
     def __init__(self, mock):
-        self.__dict__["_%s__mock" % self.__class__.__name__] = mock
+        self.__dict__["_Spy__mock"] = mock
+
+    def __getattribute__(self, item: str):
+        if item == "__dict__" or item.startswith("_Spy"):
+            return super().__getattribute__(item)
 
-    def __getattr__(self, item):
         attr = getattr(self.__mock, item)
 
-        mock_wraps = attr._mock_wraps
-        wrapped = self.__mock._mock_wraps
-        if mock_wraps is not None:
-            # Is the wrapped value present?
-            if isinstance(mock_wraps, types.MethodType):
-                # Is the wrapped value a method?
-                attr_self = mock_wraps.__self__
-                if attr_self is not self and not isinstance(attr_self, type):
-                    # If the method belongs to a Mock and method is not class method
-                    # Rebind mock method to use this Spy as self
-                    # This will allow self.method() to go back into the spy and into the Mock to be tracked
-                    attr._mock_wraps = types.MethodType(mock_wraps.__func__, self)
+        if isinstance(attr, NonCallableMock):
+            mock_wraps = attr._mock_wraps
+            wrapped = self.__mock._mock_wraps
+            if mock_wraps is not None:
+                # Is the wrapped value present?
+                if hasattr(mock_wraps, "__self__"):
+                    # Is the wrapped value a method?
+                    attr_self = mock_wraps.__self__
+                    if attr_self is not self and not isinstance(attr_self, type):
+                        # If the method belongs to a Mock and not rebound to the Spy and method is not class method
+                        if hasattr(mock_wraps, "__func__"):
+                            # If this method is not a method-wrapper
+                            # Rebind mock method to use this Spy as self
+                            # This will allow self.method() to go back into the spy and into the Mock to be tracked
+                            setattr_internal(attr, "_mock_wraps", types.MethodType(mock_wraps.__func__, self))
+                        else:
+                            # This is a method-wrapper that has no function object and has to be proxied
+                            setattr_internal(attr, "_mock_wraps",
+                                             types.MethodType(
+                                                 make_method_wrapper_closure(type(wrapped), self, mock_wraps.__name__),
+                                                 self))
+                else:
+                    # This attribute is not a wrapped method
+                    if not isinstance(mock_wraps, types.FunctionType) and hasattr(wrapped, item):
+                        # If wrapped is not a function (e.g. static method) and the underlying wrapped
+                        # has this attribute then simply return the value of that attribute directly
+                        return getattr(wrapped, item)
             else:
-                # This attribute is not a method
-                if not isinstance(mock_wraps, types.FunctionType) and hasattr(wrapped, item):
-                    # If wrapped is not a function (e.g. static method) and the underlying wrapped
-                    # has this attribute then simply return the value of that attribute directly
-                    return getattr(wrapped, item)
-        else:
-            if attr._mock_return_value is DEFAULT and hasattr(wrapped, item) and getattr(wrapped, item) is None:
-                # This attribute is not wrapped, and if it doesn't have a return value
-                # and is None then just return None
-                return None
+                if attr._mock_return_value is DEFAULT and hasattr(wrapped, item) and \
+                        getattr(wrapped, item) is None:
+                    # This attribute is not wrapped, and if it doesn't have a return value
+                    # and is None then just return None
+                    return None
 
         # In all other cases we return the attribute as we found it
         return attr
@@ -84,17 +105,74 @@ def __setattr__(self, key: str, value):
                     isinstance(mock_wraps, types.MethodType) and mock_wraps.__self__ in (self, wrapped):
                 # If attribute is not wrapped or is a method of the wrapped object and method is bound
                 # to Spy or spied mock then delegate to Mock
-                return setattr(mock, key, value)
+                return setattr_internal(mock, key, value)
 
             # Otherwise set the value directly on the object that is wrapped and is spied on
-            setattr(wrapped, key, value)
+            setattr_internal(wrapped, key, value)
         except AttributeError:
             # If Mock doesn't have this attribute delegate to Mock
-            return setattr(mock, key, value)
+            return setattr_internal(mock, key, value)
+
+
+def get_proper_attr_target(obj, key):
+    if not isinstance(obj, type) and key.startswith("__"):
+        return type(obj)
+    return obj
+
+
+def setattr_internal(obj, key, value):
+    return setattr(get_proper_attr_target(obj, key), key, value)
+
+
+def make_method_closure(method):
+    def spy_proxy_mock(self, *args, **kwargs):
+        return getattr(self, method)(*args, **kwargs)
+
+    return spy_proxy_mock
+
+
+def make_method_wrapper_closure(cls, obj, method_wrapper):
+    def method_wrapper_proxy(self, *args, **kwargs):
+        return getattr(cls, method_wrapper)(obj, *args, **kwargs)
+
+    return method_wrapper_proxy
+
+
+_magics = sorted(_all_magics)
+
+_method_wrapper_type = type(object.__str__)
+
+
+def instrument_wrapped_magic(wrapped, mock, spy):
+    with patch.dict(_calculate_return_value, {}, clear=True):
+        for magic in _magics:
+            do_wrap_magic = False
+            do_proxy_mock = False
+            if hasattr(wrapped, magic):
+                wrapped_method = getattr(wrapped, magic)
+                with patch("unittest.mock.MagicProxy.create_mock", return_value=None):
+                    mock_method = getattr(mock, magic)
+                if mock_method is None or not isinstance(mock_method, NonCallableMock):
+                    do_wrap_magic = True
+                else:
+                    do_proxy_mock = True
+                    if mock_method._mock_wraps is None and mock_method._mock_return_value is DEFAULT:
+                        do_wrap_magic = True
+
+            if do_wrap_magic:
+                do_proxy_mock = True
+                setattr_internal(mock, magic, type(mock)(wraps=wrapped_method, name=str(wrapped_method)))
+
+            if do_proxy_mock:
+                setattr_internal(spy, magic, types.MethodType(make_method_closure(magic), spy))
 
 
-def magic_spy(obj):
-    return Spy(MagicMock(spec_set=obj, wraps=obj))
+def magic_spy(wrapped, wrap_magic=True, mock_type=Mock):
+    mock = mock_type(spec_set=wrapped, wraps=wrapped)
+    spy = Spy(mock)
+    if wrap_magic:
+        instrument_wrapped_magic(wrapped, mock, spy)
+    return spy
 
 
 MagicSpy = magic_spy

src/unittest/python/mock_spy_tests.py

@@ -16,10 +16,13 @@
 # limitations under the License.
 #
 
+import warnings
 from unittest import TestCase
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, NonCallableMock
 
-from karellen.testing.mock import MagicSpy, Spy
+from karellen.testing.mock import MagicSpy, Spy, instrument_wrapped_magic
+
+warnings.simplefilter('always')
 
 
 class TestSpy(TestCase):
@@ -30,32 +33,42 @@ def __init__(self):
 
         def method_X(self):
             assert self.value == 10
-            print(self.value)
             self.method_Y()
 
         def method_Y(self):
             self.value = 9
 
         @staticmethod
         def foo():
-            print("foo")
+            pass
 
         @classmethod
         def bar(cls):
-            print("bar: " + str(cls))
+            pass
+
+        def __repr__(self):
+            return "Class_A repr"
 
     def test_magic_spy(self):
         mock = MagicSpy(TestSpy.Class_A())
 
+        mock.__repr__.assert_not_called()
+        self.assertEqual("Class_A repr", repr(mock))
+        mock.__repr__.assert_called_once_with()
+
+        self.assertIsInstance(mock, TestSpy.Class_A)
         mock.foo()
         mock.bar()
-        mock.method_X()
+        self.assertEqual("Class_A repr", str(mock))
+        mock.__str__.assert_called_once_with()
+        self.assertEqual(mock.__repr__.call_count, 2)
 
+        mock.method_X()
         mock.method_Y.assert_called_once_with()
         mock.foo.assert_called_once_with()
         mock.bar.assert_called_once_with()
 
-        self.assertEquals(mock.value, 9)
+        self.assertEqual(mock.value, 9)
         self.assertIsNone(mock.none_value)
 
     def test_spy_unwrapped(self):
@@ -69,3 +82,41 @@ def test_spy_unspecced(self):
         mock.method_X()
         mock.method_Y.assert_not_called()
         mock.bazinga = MagicMock()
+
+    def test_spy_substituted_magic_method(self):
+        wrapped = TestSpy.Class_A()
+        wrapped.__get__ = lambda *args, **kwargs: None
+
+        def mocked_repr():
+            return "Mocked repr"
+
+        def mocked_str(self):
+            return "Mocked str"
+
+        mock_proper = MagicMock(wraps=wrapped)
+        mock_proper.__str__ = mocked_str
+        mock_proper.__repr__ = MagicMock(spec_set=mocked_str, wraps=mocked_repr)
+        get_mock = MagicMock()
+        mock_proper.__get__ = get_mock
+        mock = Spy(mock_proper)
+        instrument_wrapped_magic(wrapped, mock_proper, mock)
+
+        mock.method_X()
+        mock.method_Y.assert_called_once_with()
+        self.assertEqual("Mocked str", str(mock))
+        self.assertEqual("Mocked repr", repr(mock))
+        self.assertIs(mock.__get__, get_mock)
+
+    def test_spy_no_wrapped_magic(self):
+        mock = MagicSpy(TestSpy.Class_A(), wrap_magic=False)
+        self.assertNotEqual("Class_A repr", repr(mock))
+        self.assertNotEqual("Class_A repr", str(mock))
+        self.assertNotIsInstance(mock.__str__, NonCallableMock)
+
+    def test_spy_wrapped_magic_with_magicmock(self):
+        mock = MagicSpy(TestSpy.Class_A(), mock_type=MagicMock)
+        self.assertEqual("Class_A repr", repr(mock))
+        mock.__repr__.assert_called_once_with()
+        self.assertEqual("Class_A repr", str(mock))
+        mock.__str__.assert_called_once_with()
+        self.assertIsInstance(mock.__str__, NonCallableMock)