3 files changed, 64 insertions, 17 deletions

diff --git a/UnitTests/mocktest.py b/UnitTests/mocktest.py
index 046ea16..94f86e7 100644
--- a/UnitTests/mocktest.py
+++ b/UnitTests/mocktest.py
@@ -23,6 +23,26 @@ class MockTest(TestCase):
         self.assertEquals(mock._children, {}, "children not initialised incorrectly")

         

         

+    def testSideEffect(self):

+        mock = Mock()

+        

+        test = []

+        def effect():

+            test.append(True)

+        mock.side_effect = effect

+        mock()

+        self.assertEquals(test, [True], "side effect not called")

+        self.assertTrue(mock.called, "call not recorded")

+        

+        results = [1, 2, 3]

+        def effect():

+            mock.return_value = results.pop()

+        mock.side_effect = effect

+        

+        self.assertEquals([mock(), mock(), mock()], [3, 2, 1],

+                          "side effect not used correctly")

+            

+        

     def testReset(self):

         parent = Mock()

         methods = ["something"]

@@ -70,7 +90,7 @@ class MockTest(TestCase):
         self.assertEquals(mock.call_args_list, [((sentinel.Arg,), {}), ((sentinel.Arg,), {'key': sentinel.KeyArg})], "call_args_list not set")

         

     

-    def testassert_called_with(self):

+    def testAssertCalledWith(self):

         mock = Mock()

         mock()

         

diff --git a/UnitTests/patchtest.py b/UnitTests/patchtest.py
index e7a4224..0ffb675 100644
--- a/UnitTests/patchtest.py
+++ b/UnitTests/patchtest.py
@@ -1,7 +1,7 @@
 from testcase import TestCase

 from testutils import RunTests

 

-from mock import Mock, patch, sentinel

+from mock import Mock, patch, patch_object, sentinel

 

 

 # for use in the test

@@ -11,17 +11,29 @@ something_else  = sentinel.SomethingElse
 

 class PatchTest(TestCase):

 

-    def testSinglePatch(self):

+    def testSinglePatchObject(self):

         class Something(object):

             attribute = sentinel.Original

           

         @apply

-        @patch(Something, 'attribute', sentinel.Patched)

+        @patch_object(Something, 'attribute', sentinel.Patched)

         def test():

             self.assertEquals(Something.attribute, sentinel.Patched, "unpatched")

             

         self.assertEquals(Something.attribute, sentinel.Original, "patch not restored")

 

+        

+    def testPatchObjectWithNone(self):

+        class Something(object):

+            attribute = sentinel.Original

+          

+        @apply

+        @patch_object(Something, 'attribute', None)

+        def test():

+            self.assertNone(Something.attribute, "unpatched")

+            

+        self.assertEquals(Something.attribute, sentinel.Original, "patch not restored")

+        

 

     def testMultiplePatch(self):

         class Something(object):

@@ -29,8 +41,8 @@ class PatchTest(TestCase):
             next_attribute = sentinel.Original2

         

         @apply

-        @patch(Something, 'attribute', sentinel.Patched)

-        @patch(Something, 'next_attribute', sentinel.Patched2)

+        @patch_object(Something, 'attribute', sentinel.Patched)

+        @patch_object(Something, 'next_attribute', sentinel.Patched2)

         def test():

             self.assertEquals(Something.attribute, sentinel.Patched, "unpatched")

             self.assertEquals(Something.next_attribute, sentinel.Patched2, "unpatched")

@@ -85,14 +97,14 @@ class PatchTest(TestCase):
             something2 = sentinel.Original2

         

         @apply

-        @patch(Test, 'something')

+        @patch_object(Test, 'something')

         def test(mock):

             self.assertEquals(mock, Test.something, "Mock not passed into test function")

             self.assertTrue(isinstance(mock, Mock), 

                             "patch with two arguments did not create a mock")

                             

-        @patch(Test, 'something')                    

-        @patch(Test, 'something2')

+        @patch_object(Test, 'something')                    

+        @patch_object(Test, 'something2')

         def test(this1, this2, mock1, mock2):

             self.assertEquals(this1, sentinel.this1, "Patched function didn't receive initial argument")

             self.assertEquals(this2, sentinel.this2, "Patched function didn't receive second argument")

diff --git a/mock.py b/mock.py
index 8533ddf..b3d1366 100644
--- a/mock.py
+++ b/mock.py
@@ -16,13 +16,14 @@
 __all__ = (

     'Mock',

     'patch',

+    'patch_object',

     'sentinel',

     '__version__'

 )

 

 __version__ = '0.4.0'

 

-_default_return_value = object()

+DEFAULT = object()

 

 

 class Mock(object):

@@ -38,16 +39,17 @@ class Mock(object):
         

     def reset(self):

         self.called = False

-        self._return_value = _default_return_value

+        self._return_value = DEFAULT

         self.call_args = None

         self.call_count = 0

         self.call_args_list = []

         self.method_calls = []

         self._children = {}

+        self.side_effect = None

         

     

     def __get_return_value(self):

-        if self._return_value == _default_return_value:

+        if self._return_value is DEFAULT:

             self._return_value = Mock()

         return self._return_value

     

@@ -71,7 +73,10 @@ class Mock(object):
                 break

             name = parent._name + '.' + name

             parent = parent._parent

-        

+

+        if self.side_effect is not None:

+            self.side_effect()

+            

         return self.return_value

     

     

@@ -97,9 +102,7 @@ def _importer(name):
     return mod

 

 

-def patch(target, attribute, new=None):

-    if isinstance(target, basestring):

-        target = _importer(target)

+def _patch(target, attribute, new):

         

     def patcher(func):

         original = getattr(target, attribute)

@@ -113,7 +116,7 @@ def patch(target, attribute, new=None):
         

         def patched(*args, **keywargs):

             for target, attribute, new in func.patch_list:

-                if new is None:

+                if new is DEFAULT:

                     new = Mock()

                     args += (new,)

                 setattr(target, attribute, new)

@@ -129,6 +132,18 @@ def patch(target, attribute, new=None):
     return patcher

 

 

+def patch_object(target, attribute, new=DEFAULT):

+    return _patch(target, attribute, new)

+

+

+def patch(target, attribute, new=DEFAULT):

+    if isinstance(target, basestring):

+        target = _importer(target)

+    return _patch(target, attribute, new)

+

+

+

+

 class SentinelObject(object):

     def __init__(self, name):

         self.name = name