Implemented Expect helper, which allows creating a new expect()

"function" with an explicitly provided Mocker instance.  This
helps in cases where the expression can't result in a Mock
instance (e.g. expect(iter(mock))) (#196388).

3 files changed, 35 insertions, 3 deletions

diff --git a/NEWS b/NEWS
index 5af979e..6560f17 100644
--- a/NEWS
+++ b/NEWS
@@ -8,6 +8,13 @@
   certain cases because, even though that's *not* documented, Python
   tries to use __length_hint__ in some cases.
 
+- Implemented Expect helper, which allows creating a new expect()
+  "function" with an explicitly provided Mocker instance.  This
+  helps in cases where the expression can't result in a Mock
+  instance (e.g. expect(iter(mock))) (#196388).
+
+ "function" with an explicitly provided Mocker instance.
+
 - __nonzero__ should necessarily return a boolean value, so transform Mock
   results into True (#380024).
 
diff --git a/mocker.py b/mocker.py
index c31aa20..9d8a2b6 100644
--- a/mocker.py
+++ b/mocker.py
@@ -46,7 +46,7 @@ if sys.version_info < (2, 4):
     from sets import Set as set # pragma: nocover
 
 
-__all__ = ["Mocker", "expect", "IS", "CONTAINS", "IN", "MATCH",
+__all__ = ["Mocker", "Expect", "expect", "IS", "CONTAINS", "IN", "MATCH",
            "ANY", "ARGS", "KWARGS", "MockerTestCase"]
 
 
@@ -84,6 +84,8 @@ class expect(object):
 
     """
 
+    __mocker__ = None
+
     def __init__(self, mock, attr=None):
         self._mock = mock
         self._attr = attr
@@ -92,10 +94,26 @@ class expect(object):
         return self.__class__(self._mock, attr)
 
     def __call__(self, *args, **kwargs):
-        getattr(self._mock.__mocker__, self._attr)(*args, **kwargs)
+        mocker = self.__mocker__
+        if not mocker:
+            mocker = self._mock.__mocker__
+        getattr(mocker, self._attr)(*args, **kwargs)
         return self
 
 
+def Expect(mocker):
+    """Create an expect() "function" using the given Mocker instance.
+
+    This helper allows defining an expect() "function" which works even
+    in trickier cases such as:
+
+        expect = Expect(mymocker)
+        expect(iter(mock)).generate([1, 2, 3])
+
+    """
+    return type("Expect", (expect,), {"__mocker__": mocker})
+
+
 # --------------------------------------------------------------------
 # Extensions to Python's unittest.
 
diff --git a/test.py b/test.py
index 68f3125..dc160b3 100755
--- a/test.py
+++ b/test.py
@@ -25,7 +25,7 @@ from mocker import \
     mock_returner_recorder, FunctionRunner, Orderer, SpecChecker, \
     spec_checker_recorder, match_params, ANY, IS, CONTAINS, IN, MATCH, ARGS, \
     KWARGS, MatchError, PathExecuter, ProxyReplacer, Patcher, Undefined, \
-    PatchedMethod, MockerTestCase, ReplayRestoreEvent, OnRestoreCaller
+    PatchedMethod, MockerTestCase, ReplayRestoreEvent, OnRestoreCaller, Expect
 
 
 class TestCase(unittest.TestCase):
@@ -272,6 +272,13 @@ class ExpectTest(TestCase):
         self.mocker.replay()
         self.assertEquals(obj.attr, 42)
 
+    def test_explicit_expect_instance(self):
+        obj = self.mocker.mock()
+        myexpect = Expect(self.mocker)
+        myexpect(iter(obj)).generate([1, 2, 3]).count(1, 2)
+        self.mocker.replay()
+        self.assertEquals(list(obj), [1, 2, 3])
+
 
 class MockerTestCaseTest(TestCase):