Use the more reliable expect() automatically with MockerTestCase.

2 files changed, 7 insertions, 3 deletions

diff --git a/mocker.py b/mocker.py
index 9d8a2b6..6a030e0 100644
--- a/mocker.py
+++ b/mocker.py
@@ -134,8 +134,6 @@ class MockerTestCase(unittest.TestCase):
     a few additional helper methods.
     """
 
-    expect = expect
-
     def __init__(self, methodName="runTest"):
         # So here is the trick: we take the real test method, wrap it on
         # a function that do the job we have to do, and insert it in the
@@ -184,6 +182,7 @@ class MockerTestCase(unittest.TestCase):
         self.run = run_wrapper
 
         self.mocker = Mocker()
+        self.expect = Expect(self.mocker)
 
         self.__cleanup_funcs = []
         self.__cleanup_paths = []
diff --git a/test.py b/test.py
index dc160b3..44f9efc 100755
--- a/test.py
+++ b/test.py
@@ -294,7 +294,12 @@ class MockerTestCaseTest(TestCase):
         self.assertEquals(type(self.test.mocker), Mocker)
 
     def test_has_expect(self):
-        self.assertTrue(self.test.expect is expect)
+        self.assertTrue(issubclass(self.test.expect, expect))
+
+    def test_expect_works_with_non_mocks(self):
+        # We must be using the Expect helper for this to work at all:
+        obj = self.test.mocker.mock()
+        self.test.expect(iter(obj)).generate([1,2,3])
 
     def test_attributes_are_the_same(self):
         class MyTest(MockerTestCase):