Просмотр исходного кода

Make filter test logic more obvious

Jeremy Stretch 6 лет назад
Родитель
Сommit
bc696f2e11
1 измененных файлов с 5 добавлено и 2 удалено
  1. 5 2
      netbox/extras/tests/test_filters.py

+ 5 - 2
netbox/extras/tests/test_filters.py

@@ -3,6 +3,7 @@ from django.test import TestCase
 
 from dcim.models import DeviceRole, Platform, Region, Site
 from extras.choices import *
+from extras.constants import GRAPH_MODELS
 from extras.filters import *
 from extras.models import ConfigContext, ExportTemplate, Graph
 from tenancy.models import Tenant, TenantGroup
@@ -15,7 +16,8 @@ class GraphTestCase(TestCase):
     @classmethod
     def setUpTestData(cls):
 
-        content_types = ContentType.objects.filter(model__in=['site', 'device', 'interface'])
+        # Get the first three available types
+        content_types = ContentType.objects.filter(GRAPH_MODELS)[:3]
 
         graphs = (
             Graph(name='Graph 1', type=content_types[0], template_language=TemplateLanguageChoices.LANGUAGE_DJANGO, source='http://example.com/1'),
@@ -29,7 +31,8 @@ class GraphTestCase(TestCase):
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 1)
 
     def test_type(self):
-        params = {'type': ContentType.objects.get(model='site').pk}
+        content_type = ContentType.objects.filter(GRAPH_MODELS).first()
+        params = {'type': content_type.pk}
         self.assertEqual(self.filterset(params, self.queryset).qs.count(), 1)
 
     # TODO: Remove in v2.8