Add order_by() to MockSet()

django_mock_queries/query.py

@@ -3,7 +3,7 @@
 
 from .constants import *
 from .exceptions import *
-from .utils import matches, merge, intersect
+from .utils import matches, merge, intersect, get_attribute
 
 
 class MockBase(MagicMock):
@@ -114,6 +114,18 @@ def aggregate(expr):
 
     mock_set.aggregate = MagicMock(side_effect=aggregate)
 
+    def order_by(*fields):
+        results = items
+        for field in reversed(fields):
+            is_reversed = field.startswith('-')
+            attr = field[1:] if is_reversed else field
+            results = sorted(results,
+                             key=lambda r: get_attribute(r, attr),
+                             reverse=is_reversed)
+        return MockSet(*results, cls=cls)
+
+    mock_set.order_by = MagicMock(side_effect=order_by)
+
     def latest(field):
         result = sorted(items, key=attrgetter(field), reverse=True)
         if len(result) == 0:

tests/test_query.py

@@ -188,6 +188,26 @@ def test_query_earliest_raises_error_exist_when_empty_set(self):
         self.mock_set.clear()
         self.assertRaises(Exception, self.mock_set.earliest, 'foo')
 
+    def test_query_order_by(self):
+        item_1 = MockModel(foo=1, bar='a', mock_name='item_1')
+        item_2 = MockModel(foo=1, bar='c', mock_name='item_2')
+        item_3 = MockModel(foo=2, bar='b', mock_name='item_3')
+
+        self.mock_set.add(item_1, item_3, item_2)
+        results = list(self.mock_set.order_by('foo', 'bar'))
+
+        assert results == [item_1, item_2, item_3], results
+
+    def test_query_order_by_descending(self):
+        item_1 = MockModel(foo=1, bar='c', mock_name='item_1')
+        item_2 = MockModel(foo=1, bar='a', mock_name='item_2')
+        item_3 = MockModel(foo=2, bar='b', mock_name='item_3')
+
+        self.mock_set.add(item_2, item_3, item_1)
+        results = list(self.mock_set.order_by('foo', '-bar'))
+
+        assert results == [item_1, item_2, item_3], results
+
     def test_query_implements_iterator_on_items(self):
         items = [1, 2, 3]
         assert [x for x in MockSet(*items)] == items