From 3b36a35b9a00a7f85e66f20339e8e7b872a38dd0 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@digitalocean.com>
Date: Mon, 22 Aug 2016 17:15:20 -0400
Subject: [PATCH] Fixed API tests

---
 netbox/dcim/api/serializers.py   |  2 +-
 netbox/dcim/api/views.py         | 10 +++++++++-
 netbox/dcim/tests/test_apis.py   |  9 +++++++--
 netbox/extras/api/serializers.py |  9 ++++++---
 netbox/extras/api/views.py       |  2 +-
 5 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/netbox/dcim/api/serializers.py b/netbox/dcim/api/serializers.py
index f8329d38b..583f63d8b 100644
--- a/netbox/dcim/api/serializers.py
+++ b/netbox/dcim/api/serializers.py
@@ -93,7 +93,7 @@ class RackDetailSerializer(RackSerializer):
 
     class Meta(RackSerializer.Meta):
         fields = ['id', 'name', 'facility_id', 'display_name', 'site', 'group', 'tenant', 'role', 'type', 'width',
-                  'u_height', 'comments', 'front_units', 'rear_units']
+                  'u_height', 'comments', 'custom_fields', 'front_units', 'rear_units']
 
     def get_front_units(self, obj):
         units = obj.get_rack_units(face=RACK_FACE_FRONT)
diff --git a/netbox/dcim/api/views.py b/netbox/dcim/api/views.py
index 143887624..01a8c6f61 100644
--- a/netbox/dcim/api/views.py
+++ b/netbox/dcim/api/views.py
@@ -5,6 +5,7 @@ from rest_framework.settings import api_settings
 from rest_framework.views import APIView
 
 from django.conf import settings
+from django.contrib.contenttypes.models import ContentType
 from django.http import Http404
 from django.shortcuts import get_object_or_404
 
@@ -430,6 +431,13 @@ class RelatedConnectionsView(APIView):
     Retrieve all connections related to a given console/power/interface connection
     """
 
+    def __init__(self):
+        super(RelatedConnectionsView, self).__init__()
+
+        # Custom fields
+        self.content_type = ContentType.objects.get_for_model(Device)
+        self.custom_fields = self.content_type.custom_fields.prefetch_related('choices')
+
     def get(self, request):
 
         peer_device = request.GET.get('peer-device')
@@ -454,7 +462,7 @@ class RelatedConnectionsView(APIView):
 
         # Initialize response skeleton
         response = {
-            'device': serializers.DeviceSerializer(device).data,
+            'device': serializers.DeviceSerializer(device, context={'view': self}).data,
             'console-ports': [],
             'power-ports': [],
             'interfaces': [],
diff --git a/netbox/dcim/tests/test_apis.py b/netbox/dcim/tests/test_apis.py
index b8ddf9421..f7f63c555 100644
--- a/netbox/dcim/tests/test_apis.py
+++ b/netbox/dcim/tests/test_apis.py
@@ -21,6 +21,7 @@ class SiteTest(APITestCase):
         'physical_address',
         'shipping_address',
         'comments',
+        'custom_fields',
         'count_prefixes',
         'count_vlans',
         'count_racks',
@@ -46,7 +47,8 @@ class SiteTest(APITestCase):
         'type',
         'width',
         'u_height',
-        'comments'
+        'comments',
+        'custom_fields',
     ]
 
     graph_fields = [
@@ -125,7 +127,8 @@ class RackTest(APITestCase):
         'type',
         'width',
         'u_height',
-        'comments'
+        'comments',
+        'custom_fields',
     ]
 
     detail_fields = [
@@ -141,6 +144,7 @@ class RackTest(APITestCase):
         'width',
         'u_height',
         'comments',
+        'custom_fields',
         'front_units',
         'rear_units'
     ]
@@ -337,6 +341,7 @@ class DeviceTest(APITestCase):
         'primary_ip4',
         'primary_ip6',
         'comments',
+        'custom_fields',
     ]
 
     nested_fields = ['id', 'name', 'display_name']
diff --git a/netbox/extras/api/serializers.py b/netbox/extras/api/serializers.py
index ff1e39967..4e82b4027 100644
--- a/netbox/extras/api/serializers.py
+++ b/netbox/extras/api/serializers.py
@@ -1,6 +1,6 @@
 from rest_framework import serializers
 
-from extras.models import CF_TYPE_SELECT, CustomFieldChoice, CustomFieldValue, Graph
+from extras.models import CF_TYPE_SELECT, CustomFieldChoice, Graph
 
 
 class CustomFieldSerializer(serializers.Serializer):
@@ -17,14 +17,17 @@ class CustomFieldSerializer(serializers.Serializer):
         # Attach any defined CustomFieldValues to their respective CustomFields
         for cfv in obj.custom_field_values.all():
 
-            # Suppress database lookups for CustomFieldChoices. Instead, use the cached choice set from the view
+            # Attempt to suppress database lookups for CustomFieldChoices by using the cached choice set from the view
             # context.
-            if cfv.field.type == CF_TYPE_SELECT:
+            if cfv.field.type == CF_TYPE_SELECT and hasattr(self, 'custom_field_choices'):
                 cfc = {
                     'id': int(cfv.serialized_value),
                     'value': self.context['view'].custom_field_choices[int(cfv.serialized_value)]
                 }
                 fields[cfv.field.name] = CustomFieldChoiceSerializer(instance=cfc).data
+            # Fall back to hitting the database in case we're in a view that doesn't inherit CustomFieldModelAPIView.
+            elif cfv.field.type == CF_TYPE_SELECT:
+                fields[cfv.field.name] = CustomFieldChoiceSerializer(instance=cfv.value).data
             else:
                 fields[cfv.field.name] = cfv.value
 
diff --git a/netbox/extras/api/views.py b/netbox/extras/api/views.py
index 0b0638fd7..b5928dae1 100644
--- a/netbox/extras/api/views.py
+++ b/netbox/extras/api/views.py
@@ -9,7 +9,7 @@ from django.shortcuts import get_object_or_404
 
 from circuits.models import Provider
 from dcim.models import Site, Device, Interface, InterfaceConnection
-from extras.models import CustomFieldChoice, Graph, TopologyMap, GRAPH_TYPE_INTERFACE, GRAPH_TYPE_PROVIDER, GRAPH_TYPE_SITE
+from extras.models import Graph, TopologyMap, GRAPH_TYPE_INTERFACE, GRAPH_TYPE_PROVIDER, GRAPH_TYPE_SITE
 
 from .serializers import GraphSerializer