Răsfoiți Sursa

Closes #3731: Change Graph.type to a ContentType foreign key field

Jeremy Stretch 6 ani în urmă
părinte
comite
7518174374

+ 3 - 1
docs/release-notes/version-2.7.md

@@ -126,10 +126,11 @@ PATCH) to maintain backward compatibility. This behavior will be discontinued be
 * [#3455](https://github.com/digitalocean/netbox/issues/3455) - Add tenant assignment to cluster
 * [#3455](https://github.com/digitalocean/netbox/issues/3455) - Add tenant assignment to cluster
 * [#3564](https://github.com/digitalocean/netbox/issues/3564) - Add list views for device components
 * [#3564](https://github.com/digitalocean/netbox/issues/3564) - Add list views for device components
 * [#3538](https://github.com/digitalocean/netbox/issues/3538) - Introduce a REST API endpoint for executing custom scripts
 * [#3538](https://github.com/digitalocean/netbox/issues/3538) - Introduce a REST API endpoint for executing custom scripts
+* [#3731](https://github.com/digitalocean/netbox/issues/3731) - Change Graph.type to a ContentType foreign key field
 
 
 ## API Changes
 ## API Changes
 
 
-* Choice fields now use human-friendly strings for their values instead of integers (see [#3569](https://github.com/netbox-community/netbox/issues/3569))
+* Choice fields now use human-friendly strings for their values instead of integers (see [#3569](https://github.com/netbox-community/netbox/issues/3569)).
 * Introduced `/api/extras/scripts/` endpoint for retrieving and executing custom scripts
 * Introduced `/api/extras/scripts/` endpoint for retrieving and executing custom scripts
 * dcim.ConsolePort: Added field `type`
 * dcim.ConsolePort: Added field `type`
 * dcim.ConsolePortTemplate: Added field `type`
 * dcim.ConsolePortTemplate: Added field `type`
@@ -139,4 +140,5 @@ PATCH) to maintain backward compatibility. This behavior will be discontinued be
 * dcim.PowerPortTemplate: Added field `type`
 * dcim.PowerPortTemplate: Added field `type`
 * dcim.PowerOutlet: Added field `type`
 * dcim.PowerOutlet: Added field `type`
 * dcim.PowerOutletTemplate: Added field `type`
 * dcim.PowerOutletTemplate: Added field `type`
+* extras.Graph: The `type` field has been changed to a content type foreign key. Models are specified as `<app>.<model>`; e.g. `dcim.site`.
 * virtualization.Cluster: Added field `tenant`
 * virtualization.Cluster: Added field `tenant`

+ 2 - 2
netbox/circuits/api/views.py

@@ -7,7 +7,7 @@ from circuits import filters
 from circuits.models import Provider, CircuitTermination, CircuitType, Circuit
 from circuits.models import Provider, CircuitTermination, CircuitType, Circuit
 from extras.api.serializers import RenderedGraphSerializer
 from extras.api.serializers import RenderedGraphSerializer
 from extras.api.views import CustomFieldModelViewSet
 from extras.api.views import CustomFieldModelViewSet
-from extras.models import Graph, GRAPH_TYPE_PROVIDER
+from extras.models import Graph
 from utilities.api import FieldChoicesViewSet, ModelViewSet
 from utilities.api import FieldChoicesViewSet, ModelViewSet
 from . import serializers
 from . import serializers
 
 
@@ -40,7 +40,7 @@ class ProviderViewSet(CustomFieldModelViewSet):
         A convenience method for rendering graphs for a particular provider.
         A convenience method for rendering graphs for a particular provider.
         """
         """
         provider = get_object_or_404(Provider, pk=pk)
         provider = get_object_or_404(Provider, pk=pk)
-        queryset = Graph.objects.filter(type=GRAPH_TYPE_PROVIDER)
+        queryset = Graph.objects.filter(type__model='provider')
         serializer = RenderedGraphSerializer(queryset, many=True, context={'graphed_object': provider})
         serializer = RenderedGraphSerializer(queryset, many=True, context={'graphed_object': provider})
         return Response(serializer.data)
         return Response(serializer.data)
 
 

+ 8 - 4
netbox/circuits/tests/test_api.py

@@ -1,10 +1,10 @@
+from django.contrib.contenttypes.models import ContentType
 from django.urls import reverse
 from django.urls import reverse
 from rest_framework import status
 from rest_framework import status
 
 
 from circuits.choices import *
 from circuits.choices import *
 from circuits.models import Circuit, CircuitTermination, CircuitType, Provider
 from circuits.models import Circuit, CircuitTermination, CircuitType, Provider
 from dcim.models import Site
 from dcim.models import Site
-from extras.constants import GRAPH_TYPE_PROVIDER
 from extras.models import Graph
 from extras.models import Graph
 from utilities.testing import APITestCase
 from utilities.testing import APITestCase
 
 
@@ -28,16 +28,20 @@ class ProviderTest(APITestCase):
 
 
     def test_get_provider_graphs(self):
     def test_get_provider_graphs(self):
 
 
+        provider_ct = ContentType.objects.get(app_label='circuits', model='provider')
         self.graph1 = Graph.objects.create(
         self.graph1 = Graph.objects.create(
-            type=GRAPH_TYPE_PROVIDER, name='Test Graph 1',
+            type=provider_ct,
+            name='Test Graph 1',
             source='http://example.com/graphs.py?provider={{ obj.slug }}&foo=1'
             source='http://example.com/graphs.py?provider={{ obj.slug }}&foo=1'
         )
         )
         self.graph2 = Graph.objects.create(
         self.graph2 = Graph.objects.create(
-            type=GRAPH_TYPE_PROVIDER, name='Test Graph 2',
+            type=provider_ct,
+            name='Test Graph 2',
             source='http://example.com/graphs.py?provider={{ obj.slug }}&foo=2'
             source='http://example.com/graphs.py?provider={{ obj.slug }}&foo=2'
         )
         )
         self.graph3 = Graph.objects.create(
         self.graph3 = Graph.objects.create(
-            type=GRAPH_TYPE_PROVIDER, name='Test Graph 3',
+            type=provider_ct,
+            name='Test Graph 3',
             source='http://example.com/graphs.py?provider={{ obj.slug }}&foo=3'
             source='http://example.com/graphs.py?provider={{ obj.slug }}&foo=3'
         )
         )
 
 

+ 2 - 2
netbox/circuits/views.py

@@ -6,7 +6,7 @@ from django.db.models import Count, OuterRef, Subquery
 from django.shortcuts import get_object_or_404, redirect, render
 from django.shortcuts import get_object_or_404, redirect, render
 from django.views.generic import View
 from django.views.generic import View
 
 
-from extras.models import Graph, GRAPH_TYPE_PROVIDER
+from extras.models import Graph
 from utilities.forms import ConfirmationForm
 from utilities.forms import ConfirmationForm
 from utilities.views import (
 from utilities.views import (
     BulkDeleteView, BulkEditView, BulkImportView, ObjectDeleteView, ObjectEditView, ObjectListView,
     BulkDeleteView, BulkEditView, BulkImportView, ObjectDeleteView, ObjectEditView, ObjectListView,
@@ -36,7 +36,7 @@ class ProviderView(PermissionRequiredMixin, View):
 
 
         provider = get_object_or_404(Provider, slug=slug)
         provider = get_object_or_404(Provider, slug=slug)
         circuits = Circuit.objects.filter(provider=provider).prefetch_related('type', 'tenant', 'terminations__site')
         circuits = Circuit.objects.filter(provider=provider).prefetch_related('type', 'tenant', 'terminations__site')
-        show_graphs = Graph.objects.filter(type=GRAPH_TYPE_PROVIDER).exists()
+        show_graphs = Graph.objects.filter(type__model='provider').exists()
 
 
         return render(request, 'circuits/provider.html', {
         return render(request, 'circuits/provider.html', {
             'provider': provider,
             'provider': provider,

+ 3 - 4
netbox/dcim/api/views.py

@@ -23,7 +23,6 @@ from dcim.models import (
 )
 )
 from extras.api.serializers import RenderedGraphSerializer
 from extras.api.serializers import RenderedGraphSerializer
 from extras.api.views import CustomFieldModelViewSet
 from extras.api.views import CustomFieldModelViewSet
-from extras.constants import GRAPH_TYPE_DEVICE, GRAPH_TYPE_INTERFACE, GRAPH_TYPE_SITE
 from extras.models import Graph
 from extras.models import Graph
 from ipam.models import Prefix, VLAN
 from ipam.models import Prefix, VLAN
 from utilities.api import (
 from utilities.api import (
@@ -133,7 +132,7 @@ class SiteViewSet(CustomFieldModelViewSet):
         A convenience method for rendering graphs for a particular site.
         A convenience method for rendering graphs for a particular site.
         """
         """
         site = get_object_or_404(Site, pk=pk)
         site = get_object_or_404(Site, pk=pk)
-        queryset = Graph.objects.filter(type=GRAPH_TYPE_SITE)
+        queryset = Graph.objects.filter(type__model='site')
         serializer = RenderedGraphSerializer(queryset, many=True, context={'graphed_object': site})
         serializer = RenderedGraphSerializer(queryset, many=True, context={'graphed_object': site})
         return Response(serializer.data)
         return Response(serializer.data)
 
 
@@ -357,7 +356,7 @@ class DeviceViewSet(CustomFieldModelViewSet):
         A convenience method for rendering graphs for a particular Device.
         A convenience method for rendering graphs for a particular Device.
         """
         """
         device = get_object_or_404(Device, pk=pk)
         device = get_object_or_404(Device, pk=pk)
-        queryset = Graph.objects.filter(type=GRAPH_TYPE_DEVICE)
+        queryset = Graph.objects.filter(type__model='device')
         serializer = RenderedGraphSerializer(queryset, many=True, context={'graphed_object': device})
         serializer = RenderedGraphSerializer(queryset, many=True, context={'graphed_object': device})
 
 
         return Response(serializer.data)
         return Response(serializer.data)
@@ -479,7 +478,7 @@ class InterfaceViewSet(CableTraceMixin, ModelViewSet):
         A convenience method for rendering graphs for a particular interface.
         A convenience method for rendering graphs for a particular interface.
         """
         """
         interface = get_object_or_404(Interface, pk=pk)
         interface = get_object_or_404(Interface, pk=pk)
-        queryset = Graph.objects.filter(type=GRAPH_TYPE_INTERFACE)
+        queryset = Graph.objects.filter(type__model='interface')
         serializer = RenderedGraphSerializer(queryset, many=True, context={'graphed_object': interface})
         serializer = RenderedGraphSerializer(queryset, many=True, context={'graphed_object': interface})
         return Response(serializer.data)
         return Response(serializer.data)
 
 

+ 16 - 7
netbox/dcim/tests/test_api.py

@@ -1,3 +1,4 @@
+from django.contrib.contenttypes.models import ContentType
 from django.urls import reverse
 from django.urls import reverse
 from netaddr import IPNetwork
 from netaddr import IPNetwork
 from rest_framework import status
 from rest_framework import status
@@ -12,7 +13,7 @@ from dcim.models import (
     Rack, RackGroup, RackReservation, RackRole, RearPort, Region, Site, VirtualChassis,
     Rack, RackGroup, RackReservation, RackRole, RearPort, Region, Site, VirtualChassis,
 )
 )
 from ipam.models import IPAddress, VLAN
 from ipam.models import IPAddress, VLAN
-from extras.models import Graph, GRAPH_TYPE_INTERFACE, GRAPH_TYPE_SITE
+from extras.models import Graph
 from utilities.testing import APITestCase
 from utilities.testing import APITestCase
 from virtualization.models import Cluster, ClusterType
 from virtualization.models import Cluster, ClusterType
 
 
@@ -139,16 +140,20 @@ class SiteTest(APITestCase):
 
 
     def test_get_site_graphs(self):
     def test_get_site_graphs(self):
 
 
+        site_ct = ContentType.objects.get_for_model(Site)
         self.graph1 = Graph.objects.create(
         self.graph1 = Graph.objects.create(
-            type=GRAPH_TYPE_SITE, name='Test Graph 1',
+            type=site_ct,
+            name='Test Graph 1',
             source='http://example.com/graphs.py?site={{ obj.slug }}&foo=1'
             source='http://example.com/graphs.py?site={{ obj.slug }}&foo=1'
         )
         )
         self.graph2 = Graph.objects.create(
         self.graph2 = Graph.objects.create(
-            type=GRAPH_TYPE_SITE, name='Test Graph 2',
+            type=site_ct,
+            name='Test Graph 2',
             source='http://example.com/graphs.py?site={{ obj.slug }}&foo=2'
             source='http://example.com/graphs.py?site={{ obj.slug }}&foo=2'
         )
         )
         self.graph3 = Graph.objects.create(
         self.graph3 = Graph.objects.create(
-            type=GRAPH_TYPE_SITE, name='Test Graph 3',
+            type=site_ct,
+            name='Test Graph 3',
             source='http://example.com/graphs.py?site={{ obj.slug }}&foo=3'
             source='http://example.com/graphs.py?site={{ obj.slug }}&foo=3'
         )
         )
 
 
@@ -2417,16 +2422,20 @@ class InterfaceTest(APITestCase):
 
 
     def test_get_interface_graphs(self):
     def test_get_interface_graphs(self):
 
 
+        interface_ct = ContentType.objects.get_for_model(Interface)
         self.graph1 = Graph.objects.create(
         self.graph1 = Graph.objects.create(
-            type=GRAPH_TYPE_INTERFACE, name='Test Graph 1',
+            type=interface_ct,
+            name='Test Graph 1',
             source='http://example.com/graphs.py?interface={{ obj.name }}&foo=1'
             source='http://example.com/graphs.py?interface={{ obj.name }}&foo=1'
         )
         )
         self.graph2 = Graph.objects.create(
         self.graph2 = Graph.objects.create(
-            type=GRAPH_TYPE_INTERFACE, name='Test Graph 2',
+            type=interface_ct,
+            name='Test Graph 2',
             source='http://example.com/graphs.py?interface={{ obj.name }}&foo=2'
             source='http://example.com/graphs.py?interface={{ obj.name }}&foo=2'
         )
         )
         self.graph3 = Graph.objects.create(
         self.graph3 = Graph.objects.create(
-            type=GRAPH_TYPE_INTERFACE, name='Test Graph 3',
+            type=interface_ct,
+            name='Test Graph 3',
             source='http://example.com/graphs.py?interface={{ obj.name }}&foo=3'
             source='http://example.com/graphs.py?interface={{ obj.name }}&foo=3'
         )
         )
 
 

+ 3 - 4
netbox/dcim/views.py

@@ -17,7 +17,6 @@ from django.utils.safestring import mark_safe
 from django.views.generic import View
 from django.views.generic import View
 
 
 from circuits.models import Circuit
 from circuits.models import Circuit
-from extras.constants import GRAPH_TYPE_DEVICE, GRAPH_TYPE_INTERFACE, GRAPH_TYPE_SITE
 from extras.models import Graph
 from extras.models import Graph
 from extras.views import ObjectConfigContextView
 from extras.views import ObjectConfigContextView
 from ipam.models import Prefix, VLAN
 from ipam.models import Prefix, VLAN
@@ -209,7 +208,7 @@ class SiteView(PermissionRequiredMixin, View):
             'vm_count': VirtualMachine.objects.filter(cluster__site=site).count(),
             'vm_count': VirtualMachine.objects.filter(cluster__site=site).count(),
         }
         }
         rack_groups = RackGroup.objects.filter(site=site).annotate(rack_count=Count('racks'))
         rack_groups = RackGroup.objects.filter(site=site).annotate(rack_count=Count('racks'))
-        show_graphs = Graph.objects.filter(type=GRAPH_TYPE_SITE).exists()
+        show_graphs = Graph.objects.filter(type__model='site').exists()
 
 
         return render(request, 'dcim/site.html', {
         return render(request, 'dcim/site.html', {
             'site': site,
             'site': site,
@@ -1058,8 +1057,8 @@ class DeviceView(PermissionRequiredMixin, View):
             'secrets': secrets,
             'secrets': secrets,
             'vc_members': vc_members,
             'vc_members': vc_members,
             'related_devices': related_devices,
             'related_devices': related_devices,
-            'show_graphs': Graph.objects.filter(type=GRAPH_TYPE_DEVICE).exists(),
-            'show_interface_graphs': Graph.objects.filter(type=GRAPH_TYPE_INTERFACE).exists(),
+            'show_graphs': Graph.objects.filter(type__model='device').exists(),
+            'show_interface_graphs': Graph.objects.filter(type__model='interface').exists(),
         })
         })
 
 
 
 

+ 6 - 2
netbox/extras/api/serializers.py

@@ -28,7 +28,9 @@ from .nested_serializers import *
 #
 #
 
 
 class GraphSerializer(ValidatedModelSerializer):
 class GraphSerializer(ValidatedModelSerializer):
-    type = ChoiceField(choices=GRAPH_TYPE_CHOICES)
+    type = ContentTypeField(
+        queryset=ContentType.objects.all()
+    )
 
 
     class Meta:
     class Meta:
         model = Graph
         model = Graph
@@ -38,7 +40,9 @@ class GraphSerializer(ValidatedModelSerializer):
 class RenderedGraphSerializer(serializers.ModelSerializer):
 class RenderedGraphSerializer(serializers.ModelSerializer):
     embed_url = serializers.SerializerMethodField()
     embed_url = serializers.SerializerMethodField()
     embed_link = serializers.SerializerMethodField()
     embed_link = serializers.SerializerMethodField()
-    type = ChoiceField(choices=GRAPH_TYPE_CHOICES)
+    type = ContentTypeField(
+        queryset=ContentType.objects.all()
+    )
 
 
     class Meta:
     class Meta:
         model = Graph
         model = Graph

+ 0 - 12
netbox/extras/constants.py

@@ -42,18 +42,6 @@ CUSTOMLINK_MODELS = [
     'virtualization.virtualmachine',
     'virtualization.virtualmachine',
 ]
 ]
 
 
-# Graph types
-GRAPH_TYPE_INTERFACE = 100
-GRAPH_TYPE_DEVICE = 150
-GRAPH_TYPE_PROVIDER = 200
-GRAPH_TYPE_SITE = 300
-GRAPH_TYPE_CHOICES = (
-    (GRAPH_TYPE_INTERFACE, 'Interface'),
-    (GRAPH_TYPE_DEVICE, 'Device'),
-    (GRAPH_TYPE_PROVIDER, 'Provider'),
-    (GRAPH_TYPE_SITE, 'Site'),
-)
-
 # Models which support export templates
 # Models which support export templates
 EXPORTTEMPLATE_MODELS = [
 EXPORTTEMPLATE_MODELS = [
     'circuits.circuit',
     'circuits.circuit',

+ 46 - 0
netbox/extras/migrations/0033_graph_type_to_fk.py

@@ -0,0 +1,46 @@
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+GRAPH_TYPE_CHOICES = (
+    (100, 'dcim', 'interface'),
+    (150, 'dcim', 'device'),
+    (200, 'circuits', 'provider'),
+    (300, 'dcim', 'site'),
+)
+
+
+def graph_type_to_fk(apps, schema_editor):
+    Graph = apps.get_model('extras', 'Graph')
+    ContentType = apps.get_model('contenttypes', 'ContentType')
+
+    # On a new installation (and during tests) content types might not yet exist. So, we only perform the bulk
+    # updates if a Graph has been created, which implies that we're working with a populated database.
+    if Graph.objects.exists():
+        for id, app_label, model in GRAPH_TYPE_CHOICES:
+            content_type = ContentType.objects.get(app_label=app_label, model=model)
+            Graph.objects.filter(type=id).update(type=content_type.pk)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('extras', '0032_3569_webhook_fields'),
+    ]
+
+    operations = [
+        # We have to swap the legacy IDs to ContentType PKs *before* we alter the field, to avoid triggering an
+        # IntegrityError on the ForeignKey.
+        migrations.RunPython(
+            code=graph_type_to_fk
+        ),
+        migrations.AlterField(
+            model_name='graph',
+            name='type',
+            field=models.ForeignKey(
+                limit_choices_to={'model__in': ['device', 'interface', 'provider', 'site']},
+                on_delete=django.db.models.deletion.CASCADE,
+                to='contenttypes.ContentType'
+            ),
+        ),
+    ]

+ 6 - 2
netbox/extras/models.py

@@ -408,8 +408,12 @@ class CustomLink(models.Model):
 #
 #
 
 
 class Graph(models.Model):
 class Graph(models.Model):
-    type = models.PositiveSmallIntegerField(
-        choices=GRAPH_TYPE_CHOICES
+    type = models.ForeignKey(
+        to=ContentType,
+        on_delete=models.CASCADE,
+        limit_choices_to={
+            'model__in': ['device', 'interface', 'provider', 'site']
+        }
     )
     )
     weight = models.PositiveSmallIntegerField(
     weight = models.PositiveSmallIntegerField(
         default=1000
         default=1000

+ 17 - 11
netbox/extras/tests/test_api.py

@@ -4,7 +4,6 @@ from rest_framework import status
 
 
 from dcim.models import Device, DeviceRole, DeviceType, Manufacturer, Platform, Region, Site
 from dcim.models import Device, DeviceRole, DeviceType, Manufacturer, Platform, Region, Site
 from extras.api.views import ScriptViewSet
 from extras.api.views import ScriptViewSet
-from extras.constants import GRAPH_TYPE_SITE
 from extras.models import ConfigContext, Graph, ExportTemplate, Tag
 from extras.models import ConfigContext, Graph, ExportTemplate, Tag
 from extras.scripts import BooleanVar, IntegerVar, Script, StringVar
 from extras.scripts import BooleanVar, IntegerVar, Script, StringVar
 from tenancy.models import Tenant, TenantGroup
 from tenancy.models import Tenant, TenantGroup
@@ -17,14 +16,21 @@ class GraphTest(APITestCase):
 
 
         super().setUp()
         super().setUp()
 
 
+        site_ct = ContentType.objects.get_for_model(Site)
         self.graph1 = Graph.objects.create(
         self.graph1 = Graph.objects.create(
-            type=GRAPH_TYPE_SITE, name='Test Graph 1', source='http://example.com/graphs.py?site={{ obj.name }}&foo=1'
+            type=site_ct,
+            name='Test Graph 1',
+            source='http://example.com/graphs.py?site={{ obj.name }}&foo=1'
         )
         )
         self.graph2 = Graph.objects.create(
         self.graph2 = Graph.objects.create(
-            type=GRAPH_TYPE_SITE, name='Test Graph 2', source='http://example.com/graphs.py?site={{ obj.name }}&foo=2'
+            type=site_ct,
+            name='Test Graph 2',
+            source='http://example.com/graphs.py?site={{ obj.name }}&foo=2'
         )
         )
         self.graph3 = Graph.objects.create(
         self.graph3 = Graph.objects.create(
-            type=GRAPH_TYPE_SITE, name='Test Graph 3', source='http://example.com/graphs.py?site={{ obj.name }}&foo=3'
+            type=site_ct,
+            name='Test Graph 3',
+            source='http://example.com/graphs.py?site={{ obj.name }}&foo=3'
         )
         )
 
 
     def test_get_graph(self):
     def test_get_graph(self):
@@ -44,7 +50,7 @@ class GraphTest(APITestCase):
     def test_create_graph(self):
     def test_create_graph(self):
 
 
         data = {
         data = {
-            'type': GRAPH_TYPE_SITE,
+            'type': 'dcim.site',
             'name': 'Test Graph 4',
             'name': 'Test Graph 4',
             'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=4',
             'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=4',
         }
         }
@@ -55,7 +61,7 @@ class GraphTest(APITestCase):
         self.assertHttpStatus(response, status.HTTP_201_CREATED)
         self.assertHttpStatus(response, status.HTTP_201_CREATED)
         self.assertEqual(Graph.objects.count(), 4)
         self.assertEqual(Graph.objects.count(), 4)
         graph4 = Graph.objects.get(pk=response.data['id'])
         graph4 = Graph.objects.get(pk=response.data['id'])
-        self.assertEqual(graph4.type, data['type'])
+        self.assertEqual(graph4.type, ContentType.objects.get_for_model(Site))
         self.assertEqual(graph4.name, data['name'])
         self.assertEqual(graph4.name, data['name'])
         self.assertEqual(graph4.source, data['source'])
         self.assertEqual(graph4.source, data['source'])
 
 
@@ -63,17 +69,17 @@ class GraphTest(APITestCase):
 
 
         data = [
         data = [
             {
             {
-                'type': GRAPH_TYPE_SITE,
+                'type': 'dcim.site',
                 'name': 'Test Graph 4',
                 'name': 'Test Graph 4',
                 'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=4',
                 'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=4',
             },
             },
             {
             {
-                'type': GRAPH_TYPE_SITE,
+                'type': 'dcim.site',
                 'name': 'Test Graph 5',
                 'name': 'Test Graph 5',
                 'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=5',
                 'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=5',
             },
             },
             {
             {
-                'type': GRAPH_TYPE_SITE,
+                'type': 'dcim.site',
                 'name': 'Test Graph 6',
                 'name': 'Test Graph 6',
                 'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=6',
                 'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=6',
             },
             },
@@ -91,7 +97,7 @@ class GraphTest(APITestCase):
     def test_update_graph(self):
     def test_update_graph(self):
 
 
         data = {
         data = {
-            'type': GRAPH_TYPE_SITE,
+            'type': 'dcim.site',
             'name': 'Test Graph X',
             'name': 'Test Graph X',
             'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=99',
             'source': 'http://example.com/graphs.py?site={{ obj.name }}&foo=99',
         }
         }
@@ -102,7 +108,7 @@ class GraphTest(APITestCase):
         self.assertHttpStatus(response, status.HTTP_200_OK)
         self.assertHttpStatus(response, status.HTTP_200_OK)
         self.assertEqual(Graph.objects.count(), 3)
         self.assertEqual(Graph.objects.count(), 3)
         graph1 = Graph.objects.get(pk=response.data['id'])
         graph1 = Graph.objects.get(pk=response.data['id'])
-        self.assertEqual(graph1.type, data['type'])
+        self.assertEqual(graph1.type, ContentType.objects.get_for_model(Site))
         self.assertEqual(graph1.name, data['name'])
         self.assertEqual(graph1.name, data['name'])
         self.assertEqual(graph1.source, data['source'])
         self.assertEqual(graph1.source, data['source'])