Improvements to extendmock and tests.

3 files changed, 172 insertions, 28 deletions

diff --git a/discover.py b/discover.py
index acfa7db..068b140 100644
--- a/discover.py
+++ b/discover.py
@@ -1,3 +1,4 @@
+#! /usr/bin/env python
 # Copyright Michael Foord 2009
 # Licensed under the BSD License
 # See: http://pypi.python.org/pypi/discover
diff --git a/extendmock.py b/extendmock.py
index ca7ea0e..a9a8ea9 100644
--- a/extendmock.py
+++ b/extendmock.py
@@ -3,9 +3,52 @@ from mock import Mock
 
 __all__ = ('mocksignature', 'MagicMock')
 
+"""
+extendmock adds additional features to the mock module.
+
+The mocksignature function creates wrapper functions that call a mock whilst
+preserving the signature of the original function.
+
+MagicMock is a class that adds magic method support to Mock. You add magic methods
+to mock instances in the same way you mock other attributes::
+
+    mock = Mock()
+    mock.__repr__ = lambda self: 'some string'
+
+Note that functions (or callable objects like a Mock instance) that are used for 
+magic methods must take self as the first argument.
+ 
+The only unsupported magic methods (that I'm aware of) are:
+
+* Used by Mock: __init__, __new__, __getattr__, __setattr__, __delattr__
+* Rare: __dir__,
+* Can cause bad interactions with other functionality: 
+
+    - __reversed__, __missing__, __del__, __unicode__, __getattribute__
+    - __get__, __set__, __delete__
+
+* Context manager methods are not directly supported (__enter__ and __exit__) as
+  Python doesn't lookup these methods on the class as it should (in Python 2.5 / 2.6
+  anyway).
+
+For more examples of how to use mocksignature and MagicMock see the tests.
+"""
+
+# getsignature and mocksignature heavily "inspired" by
+# the decorator module: http://pypi.python.org/pypi/decorator/
+# by Michele Simionato
+
+DELEGATE = MISSING = object()
+
 def getsignature(func, handle_defaults=False):
     assert inspect.ismethod(func) or inspect.isfunction(func)
     regargs, varargs, varkwargs, defaults = inspect.getargspec(func)
+    
+    # instance methods need to lose the self argument
+    im_self = getattr(func, 'im_self', None)
+    if im_self is not None:
+        regargs = regargs[1:]
+        
     argnames = list(regargs)
     if varargs:
         argnames.append(varargs)
@@ -32,41 +75,42 @@ def mocksignature(func, mock, handle_defaults=False):
 class MagicMock(Mock):
     pass
 
-MISSING = object()
-
-# Not supported: 
-# Used by Mock: __init__, __new__, __getattr__, __setattr__, __delattr__
-# Rare: __dir__,
-# Can cause bad interactions with other functionality: 
-#   __reversed__, __missing__, __del__, __unicode__, __getattribute__
-#   __get__, __set__, __delete__
-
-magic_methods = list((
-    "eq ne lt le gt ge" 
-    "nonzero hash getitem setitem delitem"
-    "len contains iter repr str"
-    "divmod neg pos abs invert"
-    "complex int long float oct hex index"
-    "enter exit"
-).split())
+magic_methods = [
+    ("lt le gt ge eq ne", NotImplemented),
+    ("getitem setitem delitem", TypeError),
+    ("len contains iter", TypeError),
+    ("hash repr str", DELEGATE),
+    ("nonzero", True),
+    ("divmod neg pos abs invert", TypeError),
+    ("complex int long float oct hex index", TypeError)
+]
 
-numerics = "add sub mul div truediv floordiv mod lshift rshift and xor or".split()
-magic_methods.extend(numerics)
-magic_methods.extend('i%s' % m for m in numerics)
-magic_methods.extend('r%s' % m for m in numerics)
+numerics = "add sub mul div truediv floordiv mod lshift rshift and xor or"
+inplace = ' '.join('i%s' % n for n in numerics.split())
+right = ' '.join('r%s' % n for n in numerics.split())
 
+magic_methods.extend([
+    (numerics, NotImplemented), (inplace, NotImplemented), (right, NotImplemented)
+])
 
-def get_method(name):
+def get_method(name, action):
     def func(self, *args, **kw):
-        real = getattr(self, name, MISSING)
+        real = self.__dict__.get(name, MISSING)
         if real is not MISSING:
             return real(self, *args, **kw)
-        return getattr(Mock, name)(self, *args, **kw)
+        if action is NotImplemented:
+            return NotImplemented   
+        elif isinstance(action, type) and issubclass(action, Exception):
+            raise action
+        elif action is DELEGATE:
+            return getattr(Mock, name)(self, *args, **kw)
+        return action
     func.__name__ = name
     return func
     
-for method in magic_methods:
-    name = '__%s__' % method
-    setattr(MagicMock, name, get_method(name))
+for methods, action in magic_methods:
+    for method in methods.split():
+        name = '__%s__' % method
+        setattr(MagicMock, name, get_method(name, action))
     
     
\ No newline at end of file
diff --git a/tests/testextendmock.py b/tests/testextendmock.py
index 5b569b9..7307de7 100644
--- a/tests/testextendmock.py
+++ b/tests/testextendmock.py
@@ -1,6 +1,7 @@
 # Copyright (C) 2007-20010 Michael Foord

 # E-mail: fuzzyman AT voidspace DOT org DOT uk

 # http://www.voidspace.org.uk/python/mock/

+from __future__ import with_statement

 

 import os

 import sys

@@ -48,8 +49,106 @@ class TestMockSignature(TestCase):
         self.assertEquals(f.method('foo', 'bar'), 3)

         mock.assert_called_with('foo', 'bar')

         

-        

 

+class TestMagicMock(TestCase):

+    

+    def testRepr(self):

+        mock = MagicMock()

+        self.assertEqual(repr(mock), object.__repr__(mock))

+        mock.__repr__ = lambda self: 'foo'

+        self.assertEqual(repr(mock), 'foo')

+    

+    

+    def testDictMethods(self):

+        mock = MagicMock()

+        

+        self.assertRaises(TypeError, lambda: mock['foo'])

+        def _del():

+            del mock['foo']

+        def _set():

+            mock['foo'] = 3

+        self.assertRaises(TypeError, _del)

+        self.assertRaises(TypeError, _set)

+        

+        _dict = {}

+        def getitem(s, name):

+            return _dict[name]    

+        def setitem(s, name, value):

+            _dict[name] = value

+        def delitem(s, name):

+            del _dict[name]

+        

+        mock.__setitem__ = setitem

+        mock.__getitem__ = getitem

+        mock.__delitem__ = delitem

+        

+        self.assertRaises(KeyError, lambda: mock['foo'])

+        mock['foo'] = 'bar'

+        self.assertEquals(_dict, {'foo': 'bar'})

+        self.assertEquals(mock['foo'], 'bar')

+        del mock['foo']

+        self.assertEquals(_dict, {})

+            

+            

+    def testNumeric(self):

+        original = mock = MagicMock()

+        mock.value = 0

+        

+        self.assertEqual(mock.__add__(3), NotImplemented)

+        

+        def add(self, other):

+            mock.value += other

+            return self

+        mock.__add__ = add

+        self.assertEqual(mock + 3, mock)

+        self.assertEqual(mock.value, 3)

+        

+        self.assertEqual(mock.__iadd__(3), NotImplemented)        

+        mock.__iadd__ = add

+        mock += 6

+        self.assertEqual(mock, original)

+        self.assertEqual(mock.value, 9)

+        

+        self.assertEqual(mock.__radd__(3), NotImplemented)

+        mock.__radd__ = add

+        self.assertEqual(7 + mock, mock)

+        self.assertEqual(mock.value, 16)

+    

+    def testHash(self):

+        mock = MagicMock()

+        # test delegation

+        self.assertEqual(hash(mock), Mock.__hash__(mock))

+        

+        def _hash(s):

+            return 3

+        mock.__hash__ = _hash

+        self.assertEqual(hash(mock), 3)

+    

+    def testNonZero(self):

+        m = MagicMock()

+        self.assertTrue(bool(m))

+        

+        nonzero = lambda s: False

+        m.__nonzero__ = nonzero

+        self.assertFalse(bool(m))

+        

+    def testComparison(self):

+        self. assertEqual(MagicMock() < 3, object() < 3)

+        self. assertEqual(MagicMock() > 3, object() > 3)

+        self. assertEqual(MagicMock() <= 3, object() <= 3)

+        self. assertEqual(MagicMock() >= 3, object() >= 3)

+        

+        mock = MagicMock()

+        def comp(s, o):

+            print 'foo'

+            return True

+        mock.__lt__ = mock.__gt__ = mock.__le__ = mock.__ge__ = comp

+        self. assertTrue(MagicMock() < 3)

+        self. assertTrue(MagicMock() > 3)

+        self. assertTrue(MagicMock() <= 3)

+        self. assertTrue(MagicMock() >= 3)

+        

+        

 if __name__ == '__main__':

     unittest.main()