Support query set returning methods

.travis.yml

@@ -6,8 +6,8 @@ python:
 env:
   - DJANGO=1.2.7
   - DJANGO=1.3.1
-  - DJANGO=1.4.5
-  - DJANGO=1.5.1
+  - DJANGO=1.4.6
+  - DJANGO=1.5.2
 install:
   - sudo apt-get install libevent-dev
   - sudo pip install Django==$DJANGO --use-mirrors

README.rst

@@ -11,12 +11,13 @@ Inside your virtualenv::
    >>> settings.configure() # required to convince Django it's properly configured
    >>> from mock_django.query import QuerySetMock
    >>> class Post(object): pass
+   ...
    >>> qs = QuerySetMock(Post, 1, 2, 3)
    >>> list(qs.all())
    [1, 2, 3]
    >>> qs.count()
    3
-   >>> qs.all().filter()
+   >>> list(qs.all().filter())
    [1, 2, 3]
 
 See tests for more examples.

mock_django/query.py

@@ -5,23 +5,36 @@
 :copyright: (c) 2012 DISQUS.
 :license: Apache License 2.0, see LICENSE for more details.
 """
+import copy
 
 import mock
 from .shared import SharedMock
 
 __all__ = ('QuerySetMock',)
 
+QUERYSET_RETURNING_METHODS = ['filter', 'exclude', 'order_by', 'reverse',
+                              'distinct', 'none', 'all', 'select_related',
+                              'prefetch_related', 'defer', 'only', 'using',
+                              'select_for_update']
+
 
 def QuerySetMock(model, *return_value):
     """
+    Get a SharedMock that returns self for most attributes and a new copy of
+    itself for any method that ordinarily generates QuerySets.
+
     Set the results to two items:
 
+    >>> class Post(object): pass
     >>> objects = QuerySetMock(Post, 'return', 'values')
     >>> assert list(objects.filter()) == list(objects.all())
 
     Force an exception:
 
     >>> objects = QuerySetMock(Post, Exception())
+
+    Chain calls:
+    >>> objects.all().filter(filter_arg='dummy')
     """
 
     def make_get(self, model):
@@ -35,6 +48,11 @@ def _get(*a, **k):
                 raise model.DoesNotExist
         return _get
 
+    def make_qs_returning_method(self):
+        def _qs_returning_method(*a, **k):
+            return copy.deepcopy(self)
+        return _qs_returning_method
+
     def make_getitem(self):
         def _getitem(k):
             if isinstance(k, slice):
@@ -62,18 +80,22 @@ def _iterator(*a, **k):
     else:
         model = mock.MagicMock()
 
-    m = SharedMock(reserved=['count', 'exists'])
+    m = SharedMock(reserved=['count', 'exists'] + QUERYSET_RETURNING_METHODS)
     m.__start = None
     m.__stop = None
     m.__iter__.side_effect = lambda: iter(m.iterator())
     m.__getitem__.side_effect = make_getitem(m)
-    m.model = model
-    m.get = make_get(m, actual_model)
     m.__nonzero__.side_effect = lambda: bool(return_value)
     m.__len__.side_effect = lambda: len(return_value)
     m.exists.side_effect = m.__nonzero__
     m.count.side_effect = m.__len__
 
+    m.model = model
+    m.get = make_get(m, actual_model)
+
+    for method_name in QUERYSET_RETURNING_METHODS:
+        setattr(m, method_name, make_qs_returning_method(m))
+
     # Note since this is a SharedMock, *all* auto-generated child
     # attributes will have the same side_effect ... might not make
     # sense for some like count().

mock_django/shared.py

@@ -70,8 +70,7 @@ def assert_chain_calls(self, *calls):
                 not_found.append(kall)
         if not_found:
             if self.__parent.mock_calls:
-                message = '%r not all found in call list, %d other(s) were:\n%r' % (not_found, len(self.__parent.mock_calls),
-                    self.__parent.mock_calls)
+                message = '%r not all found in call list, %d other(s) were:\n%r' % (not_found, len(self.__parent.mock_calls), self.__parent.mock_calls)
             else:
                 message = 'no calls were found'
 

tests/mock_django/http/tests.py

@@ -19,7 +19,7 @@ def test_instance(self):
         self.assertEqual({}, wsgi_r.META)
         self.assertEqual({}, wsgi_r.GET)
         self.assertEqual({}, wsgi_r.POST)
-        
+
     def test__get_request(self):
         wsgi_r = WsgiHttpRequest()
         expected_items = MergeDict({}, {}).items()

tests/mock_django/query/tests.py

@@ -1,16 +1,19 @@
 from mock_django.query import QuerySetMock
 from unittest2 import TestCase
 
+
 class TestException(Exception):
     pass
 
+
 class TestModel(object):
     def foo(self):
         pass
 
     def bar(self):
         return 'bar'
 
+
 class QuerySetTestCase(TestCase):
     def test_vals_returned(self):
         qs = QuerySetMock(None, 1, 2, 3)
@@ -52,6 +55,23 @@ def test_raises_exception_when_accessed(self):
         qs = QuerySetMock(None, TestException())
         self.assertRaises(TestException, lambda x: x[0], qs)
 
+    def test_chaining_calls_works(self):
+        """
+        Chained calls to QS-returning methods should return new QuerySetMocks.
+        """
+        qs = QuerySetMock(None, 1, 2, 3)
+        qs.all().filter(filter_arg='dummy')
+        qs.filter(filter_arg='dummy').order_by('-date')
+
+    def test_chained_calls_return_new_querysetmocks(self):
+        qs = QuerySetMock(None, 1, 2, 3)
+        qs_all = qs.all()
+        qs_filter = qs.filter()
+        qs_all_filter = qs.filter().all()
+
+        self.assertIsNot(qs_all, qs_filter)
+        self.assertIsNot(qs_filter, qs_all_filter)
+
     # Test reserved methods
     def test_count_is_scalar(self):
         qs = QuerySetMock(None, 1, 2, 3)

tests/mock_django/signals/tests.py

@@ -11,6 +11,7 @@ def test_mock_receiver(self):
             self.assertEqual(receiver.call_count, 1)
 
         sentinel = {}
+
         def side_effect(*args, **kwargs):
             return sentinel