Added support for __contains__, thanks to Adam Lowry.

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

2 files changed, 71 insertions, 0 deletions

diff --git a/mox.py b/mox.py
index 0dd4cc2..2fddc0b 100755
--- a/mox.py
+++ b/mox.py
@@ -484,6 +484,33 @@ class MockObject(MockAnything, object):
     # Otherwise, create a mock method __getitem__.
     return self._CreateMockMethod('__getitem__')(key)
 
+  def __contains__(self, key):
+    """Provide custom logic for mocking classes that contain items.
+
+    Args:
+      key: Key to look in container for.
+
+    Returns:
+      Expected return value in replay mode.  A MockMethod object for the
+      __contains__ method that has already been called if not in replay mode.
+
+    Raises:
+      TypeError if the underlying class does not implement __contains__
+      UnexpectedMethodCaller if the object does not expect the call to
+      __contains__.
+
+    """
+    contains = self._class_to_mock.__dict__.get('__contains__', None)
+
+    if contains is None:
+      raise TypeError('unsubscriptable object')
+
+    if self._replay_mode:
+      return MockMethod('__contains__', self._expected_calls_queue,
+                        self._replay_mode)(key)
+
+    return self._CreateMockMethod('__contains__')(key)
+
   def __call__(self, *params, **named_params):
     """Provide custom logic for mocking classes that are callable."""
 
diff --git a/mox_test.py b/mox_test.py
index 2d8e435..b626889 100755
--- a/mox_test.py
+++ b/mox_test.py
@@ -748,6 +748,46 @@ class MockObjectTest(unittest.TestCase):
 
     dummy._Verify()
 
+  def testMockContains_ExpectedContains_Success(self):
+    """Test that __contains__ gets mocked in Dummy.
+
+    In this test, _Verify() succeeds.
+    """
+    dummy = mox.MockObject(TestClass)
+    dummy.__contains__('X').AndReturn(True)
+
+    dummy._Replay()
+
+    self.failUnless('X' in dummy)
+
+    dummy._Verify()
+
+  def testMockContains_ExpectedContains_NoSuccess(self):
+    """Test that __contains__() gets mocked in Dummy.
+
+    In this test, _Verify() fails.
+    """
+    dummy = mox.MockObject(TestClass)
+    dummy.__contains__('X').AndReturn('True')
+
+    dummy._Replay()
+
+    # NOT doing 'X' in dummy
+
+    self.assertRaises(mox.ExpectedMethodCallsError, dummy._Verify)
+
+  def testMockContains_ExpectedContains_NonmatchingParameter(self):
+    """Test that __contains__ fails if other parameters are expected."""
+    dummy = mox.MockObject(TestClass)
+    dummy.__contains__('X').AndReturn(True)
+
+    dummy._Replay()
+
+    def call(): return 'Y' in dummy
+
+    self.assertRaises(mox.UnexpectedMethodCallError, call)
+
+    dummy._Verify()
 
 class MoxTest(unittest.TestCase):
   """Verify Mox works correctly."""
@@ -1338,6 +1378,10 @@ class TestClass:
     """Set the value for key to value."""
     self.d[key] = value
 
+  def __contains__(self, key):
+     """Returns True if d contains the key."""
+     return key in self.d
+
 
 class ChildClass(TestClass):
   """This inherits from TestClass."""