Add tests for manager methods & utility functions

netbox/core/tests/test_models.py

@@ -1,7 +1,10 @@
+from django.contrib.contenttypes.models import ContentType
+from django.core.exceptions import ObjectDoesNotExist
 from django.test import TestCase
 
-from core.models import DataSource
+from core.models import DataSource, ObjectType
 from core.choices import ObjectChangeActionChoices
+from dcim.models import Site, Location, Device
 from netbox.constants import CENSOR_TOKEN, CENSOR_TOKEN_CHANGED
 
 
@@ -120,3 +123,80 @@ class DataSourceChangeLoggingTestCase(TestCase):
         self.assertEqual(objectchange.prechange_data['parameters']['password'], CENSOR_TOKEN)
         self.assertEqual(objectchange.postchange_data['parameters']['username'], 'username2')
         self.assertEqual(objectchange.postchange_data['parameters']['password'], CENSOR_TOKEN)
+
+
+class ObjectTypeTest(TestCase):
+
+    def test_create(self):
+        """
+        Test that an ObjectType created for a given app_label & model name will be automatically assigned to
+        the appropriate ContentType.
+        """
+        kwargs = {
+            'app_label': 'foo',
+            'model': 'bar',
+        }
+        ct = ContentType.objects.create(**kwargs)
+        ot = ObjectType.objects.create(**kwargs)
+        self.assertEqual(ot.contenttype_ptr, ct)
+
+    def test_get_by_natural_key(self):
+        """
+        Test that get_by_natural_key() returns the appropriate ObjectType.
+        """
+        self.assertEqual(
+            ObjectType.objects.get_by_natural_key('dcim', 'site'),
+            ObjectType.objects.get(app_label='dcim', model='site')
+        )
+        with self.assertRaises(ObjectDoesNotExist):
+            ObjectType.objects.get_by_natural_key('foo', 'bar')
+
+    def test_get_for_id(self):
+        """
+        Test that get_by_id() returns the appropriate ObjectType.
+        """
+        ot = ObjectType.objects.get_by_natural_key('dcim', 'site')
+        self.assertEqual(
+            ObjectType.objects.get_for_id(ot.pk),
+            ObjectType.objects.get(pk=ot.pk)
+        )
+        with self.assertRaises(ObjectDoesNotExist):
+            ObjectType.objects.get_for_id(0)
+
+    def test_get_for_model(self):
+        """
+        Test that get_by_model() returns the appropriate ObjectType.
+        """
+        self.assertEqual(
+            ObjectType.objects.get_for_model(Site),
+            ObjectType.objects.get_by_natural_key('dcim', 'site')
+        )
+
+    def test_get_for_models(self):
+        """
+        Test that get_by_models() returns the appropriate ObjectType mapping.
+        """
+        self.assertEqual(
+            ObjectType.objects.get_for_models(Site, Location, Device),
+            {
+                Site: ObjectType.objects.get_by_natural_key('dcim', 'site'),
+                Location: ObjectType.objects.get_by_natural_key('dcim', 'location'),
+                Device: ObjectType.objects.get_by_natural_key('dcim', 'device'),
+            }
+        )
+
+    def test_public(self):
+        """
+        Test that public() returns only ObjectTypes for public models.
+        """
+        public_ots = ObjectType.objects.public()
+        self.assertIn(ObjectType.objects.get_by_natural_key('dcim', 'site'), public_ots)
+        self.assertNotIn(ObjectType.objects.get_by_natural_key('extras', 'taggeditem'), public_ots)
+
+    def test_with_feature(self):
+        """
+        Test that with_feature() returns only ObjectTypes for models which support the specified feature.
+        """
+        bookmarks_ots = ObjectType.objects.with_feature('bookmarks')
+        self.assertIn(ObjectType.objects.get_by_natural_key('dcim', 'site'), bookmarks_ots)
+        self.assertNotIn(ObjectType.objects.get_by_natural_key('dcim', 'cabletermination'), bookmarks_ots)

netbox/netbox/tests/test_model_features.py

@@ -0,0 +1,43 @@
+from django.test import TestCase
+
+from core.models import AutoSyncRecord, DataSource
+from extras.models import CustomLink
+from netbox.models.features import get_model_features, has_feature, model_is_public
+
+
+class ModelFeaturesTestCase(TestCase):
+
+    def test_model_is_public(self):
+        """
+        Test that the is_public() utility function returns True for public models only.
+        """
+        # Public model
+        self.assertFalse(hasattr(DataSource, '_netbox_private'))
+        self.assertTrue(model_is_public(DataSource))
+
+        # Private model
+        self.assertTrue(getattr(AutoSyncRecord, '_netbox_private'))
+        self.assertFalse(model_is_public(AutoSyncRecord))
+
+    def test_has_feature(self):
+        """
+        Test the functionality of the has_feature() utility function.
+        """
+        # Sanity checking
+        self.assertTrue(hasattr(DataSource, 'bookmarks'), "Invalid test?")
+        self.assertFalse(hasattr(AutoSyncRecord, 'bookmarks'), "Invalid test?")
+
+        self.assertTrue(has_feature(DataSource, 'bookmarks'))
+        self.assertFalse(has_feature(AutoSyncRecord, 'bookmarks'))
+
+    def test_get_model_features(self):
+        """
+        Check that get_model_features() returns the expected features for a model.
+        """
+        # Sanity checking
+        self.assertTrue(hasattr(CustomLink, 'clone'), "Invalid test?")
+        self.assertFalse(hasattr(CustomLink, 'bookmarks'), "Invalid test?")
+
+        features = get_model_features(CustomLink)
+        self.assertIn('cloning', features)
+        self.assertNotIn('bookmarks', features)