Fix a bug where the interface is not enforced when stubbing out a first order method,

or on the __call__ function of a callable object.



git-svn-id: http://pymox.googlecode.com/svn/trunk@43 b1010a0a-674b-0410-b734-77272b80c875

3 files changed, 33 insertions, 6 deletions

diff --git a/mox.py b/mox.py
index 3f78c9e..f88e098 100755
--- a/mox.py
+++ b/mox.py
@@ -174,8 +174,8 @@ class Mox(object):
 
   # A list of types that should be stubbed out with MockObjects (as
   # opposed to MockAnythings).
-  _USE_MOCK_OBJECT = [types.ClassType, types.InstanceType, types.ModuleType,
-                      types.ObjectType, types.TypeType]
+  _USE_MOCK_OBJECT = [types.ClassType, types.FunctionType, types.InstanceType,
+                      types.ModuleType, types.ObjectType, types.TypeType]
 
   def __init__(self):
     """Initialize a new Mox."""
@@ -246,13 +246,13 @@ class Mox(object):
     """
 
     attr_to_replace = getattr(obj, attr_name)
+    attr_type = type(attr_to_replace)
 
-    # Check for a MockAnything. This could cause confusing problems later on.
-    if attr_to_replace == MockAnything():
+    if attr_type == MockAnything or attr_type == MockObject:
       raise TypeError('Cannot mock a MockAnything! Did you remember to '
                       'call UnsetStubs in your previous test?')
 
-    if type(attr_to_replace) in self._USE_MOCK_OBJECT and not use_mock_anything:
+    if attr_type in self._USE_MOCK_OBJECT and not use_mock_anything:
       stub = self.CreateMock(attr_to_replace)
     else:
       stub = self.CreateMockAnything(description='Stub for %s' % attr_to_replace)
@@ -640,7 +640,16 @@ class MockObject(MockAnything, object):
 
     # Because the call is happening directly on this object instead of a method,
     # the call on the mock method is made right here
-    mock_method = self._CreateMockMethod('__call__')
+
+    # If we are mocking a Function, then use the function, and not the
+    # __call__ method
+    method = None
+    if type(self._class_to_mock) == types.FunctionType:
+      method = self._class_to_mock;
+    else:
+      method = getattr(self._class_to_mock, '__call__')
+    mock_method = self._CreateMockMethod('__call__', method_to_mock=method)
+
     return mock_method(*params, **named_params)
 
   @property
diff --git a/mox_test.py b/mox_test.py
index c7432df..9f88a9e 100755
--- a/mox_test.py
+++ b/mox_test.py
@@ -1176,6 +1176,11 @@ class MoxTest(unittest.TestCase):
 
     self.assertRaises(mox.UnexpectedMethodCallError, mock_obj, "ZOOBAZ")
 
+  def testCallableObjectVerifiesSignature(self):
+    mock_obj = self.mox.CreateMock(CallableClass)
+    # Too many arguments
+    self.assertRaises(AttributeError, mock_obj, "foo", "bar")
+
   def testUnorderedGroup(self):
     """Test that using one unordered group works."""
     mock_obj = self.mox.CreateMockAnything()
@@ -1507,6 +1512,13 @@ class MoxTest(unittest.TestCase):
     self.assertRaises(TypeError, self.mox.StubOutWithMock, TestClass,
                       'MyStaticMethod')
 
+  def testStubOutFirstClassMethodVerifiesSignature(self):
+    self.mox.StubOutWithMock(mox_test_helper, 'MyTestFunction')
+
+    # Wrong number of arguments
+    self.assertRaises(AttributeError, mox_test_helper.MyTestFunction, 1)
+    self.mox.UnsetStubs()
+
   def testStubOutObject(self):
     """Test than object is replaced with a Mock."""
 
@@ -1784,6 +1796,9 @@ class TestClass:
   def ValidCall(self):
     pass
 
+  def MethodWithArgs(self, one, two, nine=None):
+    pass
+
   def OtherValidCall(self):
     pass
 
diff --git a/mox_test_helper.py b/mox_test_helper.py
index b4bfdec..2d04674 100755
--- a/mox_test_helper.py
+++ b/mox_test_helper.py
@@ -93,3 +93,6 @@ class TestClassFromAnotherModule(object):
 
   def Value():
     return "Not mock"
+
+def MyTestFunction(one, two, nine=None):
+  pass