Merge pull request #74 from stphivos/refactor/remove_duplication

Refactor more maintainability issues on duplication

django_mock_queries/mocks.py

@@ -100,12 +100,24 @@ def compiler(queryset, connection, using, **kwargs):
     Model.refresh_from_db = Mock()  # Make this into a noop.
 
 
-class MockOneToManyMap(object):
+class MockMap(object):
     def __init__(self, original):
         """ Wrap a mock mapping around the original one-to-many relation. """
         self.map = {}
         self.original = original
 
+    def __set__(self, instance, value):
+        """ Set a related object for an instance. """
+
+        self.map[id(instance)] = (weakref.ref(instance), value)
+
+    def __getattr__(self, name):
+        """ Delegate all other calls to the original. """
+
+        return getattr(self.original, name)
+
+
+class MockOneToManyMap(MockMap):
     def __get__(self, instance, owner):
         """ Look in the map to see if there is a related set.
 
@@ -129,23 +141,8 @@ def __get__(self, instance, owner):
 
         return related_objects
 
-    def __set__(self, instance, value):
-        """ Set a related object for an instance. """
-
-        self.map[id(instance)] = (weakref.ref(instance), value)
-
-    def __getattr__(self, name):
-        """ Delegate all other calls to the original. """
-
-        return getattr(self.original, name)
-
-
-class MockOneToOneMap(object):
-    def __init__(self, original):
-        """ Wrap a mock mapping around the original one-to-one relation. """
-        self.map = {}
-        self.original = original
 
+class MockOneToOneMap(MockMap):
     def __get__(self, instance, owner):
         """ Look in the map to see if there is a related object.
 
@@ -170,16 +167,6 @@ def __get__(self, instance, owner):
             )
         return related_object
 
-    def __set__(self, instance, value):
-        """ Set a related object for an instance. """
-
-        self.map[id(instance)] = (weakref.ref(instance), value)
-
-    def __getattr__(self, name):
-        """ Delegate all other calls to the original. """
-
-        return getattr(self.original, name)
-
 
 def find_all_models(models):
     """ Yield all models and their parents. """

examples/users/analytics/api.py

@@ -13,3 +13,6 @@ def create_user(self, **attrs):
     def today_visitors_count(self):
         result = User.objects.filter(last_login__gte=date.today()).aggregate(Count('last_login'))
         return result['last_login__count']
+
+    def staff_usernames(self):
+        return User.objects.filter(is_staff=True).values_list('username', flat=True)

examples/users/analytics/tests.py

@@ -24,13 +24,17 @@
 
 @skipIfDBFeature('is_mocked')
 class TestApi(TestCase):
-    def test_api_create_user(self):
-        start_count = User.objects.count()
+    def setUp(self):
+        self.api = AnalyticsApi()
 
-        User.objects.create(username='bob')
-        final_count = User.objects.count()
+    def test_api_create_user(self):
+        _ = User.objects.create(username='plain1')
+        _ = User.objects.create(username='plain2')
+        staff1 = User.objects.create(username='staff1', is_staff=True)
+        staff2 = User.objects.create(username='staff2', is_staff=True)
 
-        self.assertEqual(start_count + 1, final_count)
+        usernames = [str(x) for x in self.api.staff_usernames()]
+        self.assertEqual(usernames, [staff1.username, staff2.username])
 
 
 @mocked_relations(User)