diff --git a/netbox/core/tests/test_models.py b/netbox/core/tests/test_models.py index ff71c2e88..28225c7a6 100644 --- a/netbox/core/tests/test_models.py +++ b/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) diff --git a/netbox/netbox/tests/test_model_features.py b/netbox/netbox/tests/test_model_features.py new file mode 100644 index 000000000..13934bff6 --- /dev/null +++ b/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)