Add test for CustomFieldManager.get_for_model()

netbox/extras/models/customfields.py

@@ -62,13 +62,8 @@ class CustomFieldManager(models.Manager):
         """
         Return all CustomFields assigned to the given model.
         """
-        model = model._meta.concrete_model
+        content_type = ContentType.objects.get_for_model(model._meta.concrete_model)
-
+        return self.get_queryset().filter(obj_type=content_type)
-        # Fetch from the database if the model's CustomFields have not been cached
-        content_type = ContentType.objects.get_for_model(model)
-        customfields = CustomField.objects.filter(obj_type=content_type)
-
-        return customfields
 
 
 class CustomField(models.Model):

netbox/extras/tests/test_customfields.py

@@ -99,6 +99,19 @@ class CustomFieldTest(TestCase):
         cf.delete()
 
 
+class CustomFieldManagerTest(TestCase):
+
+    def setUp(self):
+        content_type = ContentType.objects.get_for_model(Site)
+        custom_field = CustomField(type=CustomFieldTypeChoices.TYPE_TEXT, name='text_field', default='foo')
+        custom_field.save()
+        custom_field.obj_type.set([content_type])
+
+    def test_get_for_model(self):
+        self.assertEqual(CustomField.objects.get_for_model(Site).count(), 1)
+        self.assertEqual(CustomField.objects.get_for_model(VirtualMachine).count(), 0)
+
+
 class CustomFieldAPITest(APITestCase):
 
     @classmethod