#2350 - added virtual machines and test cases

netbox/dcim/tests/test_api.py

@@ -1791,6 +1791,16 @@ def setUp(self):
             site=self.site1,
             cluster=self.cluster1
         )
+        self.device_with_context_data = Device.objects.create(
+            device_type=self.devicetype1,
+            device_role=self.devicerole1,
+            name='Device with context data',
+            site=self.site1,
+            local_context_data={
+                'A': 1,
+                'B': 2
+            }
+        )
 
     def test_get_device(self):
 
@@ -1806,7 +1816,7 @@ def test_list_devices(self):
         url = reverse('dcim-api:device-list')
         response = self.client.get(url, **self.header)
 
-        self.assertEqual(response.data['count'], 3)
+        self.assertEqual(response.data['count'], 4)
 
     def test_list_devices_brief(self):
 
@@ -1832,7 +1842,7 @@ def test_create_device(self):
         response = self.client.post(url, data, format='json', **self.header)
 
         self.assertHttpStatus(response, status.HTTP_201_CREATED)
-        self.assertEqual(Device.objects.count(), 4)
+        self.assertEqual(Device.objects.count(), 5)
         device4 = Device.objects.get(pk=response.data['id'])
         self.assertEqual(device4.device_type_id, data['device_type'])
         self.assertEqual(device4.device_role_id, data['device_role'])
@@ -1867,7 +1877,7 @@ def test_create_device_bulk(self):
         response = self.client.post(url, data, format='json', **self.header)
 
         self.assertHttpStatus(response, status.HTTP_201_CREATED)
-        self.assertEqual(Device.objects.count(), 6)
+        self.assertEqual(Device.objects.count(), 7)
         self.assertEqual(response.data[0]['name'], data[0]['name'])
         self.assertEqual(response.data[1]['name'], data[1]['name'])
         self.assertEqual(response.data[2]['name'], data[2]['name'])
@@ -1891,7 +1901,7 @@ def test_update_device(self):
         response = self.client.put(url, data, format='json', **self.header)
 
         self.assertHttpStatus(response, status.HTTP_200_OK)
-        self.assertEqual(Device.objects.count(), 3)
+        self.assertEqual(Device.objects.count(), 4)
         device1 = Device.objects.get(pk=response.data['id'])
         self.assertEqual(device1.device_type_id, data['device_type'])
         self.assertEqual(device1.device_role_id, data['device_role'])
@@ -1906,7 +1916,21 @@ def test_delete_device(self):
         response = self.client.delete(url, **self.header)
 
         self.assertHttpStatus(response, status.HTTP_204_NO_CONTENT)
-        self.assertEqual(Device.objects.count(), 2)
+        self.assertEqual(Device.objects.count(), 3)
+
+    def test_config_context_included_by_default_in_list_view(self):
+
+        url = reverse('dcim-api:device-list') + '?slug=device-with-context-data'
+        response = self.client.get(url, **self.header)
+
+        self.assertEqual(response.data['results'][0].get('config_context', {}).get('A'), 1)
+
+    def test_config_context_excluded(self):
+
+        url = reverse('dcim-api:device-list') + '?exclude=config_context'
+        response = self.client.get(url, **self.header)
+
+        self.assertFalse('config_context' in response.data['results'][0])
 
 
 class ConsolePortTest(APITestCase):

netbox/virtualization/api/views.py

@@ -50,16 +50,23 @@ class VirtualMachineViewSet(CustomFieldModelViewSet):
 
     def get_serializer_class(self):
         """
-        Include rendered config context when retrieving a single VirtualMachine.
+        Select the specific serializer based on the request context.
+
+        If the `brief` query param equates to True, return the NestedVirtualMachineSerializer
+
+        If the `exclude` query param includes `config_context` as a value, return the VirtualMachineSerializer
+
+        Else, return the VirtualMachineWithConfigContextSerializer
         """
-        if self.action == 'retrieve':
-            return serializers.VirtualMachineWithConfigContextSerializer
 
         request = self.get_serializer_context()['request']
         if request.query_params.get('brief', False):
             return serializers.NestedVirtualMachineSerializer
 
-        return serializers.VirtualMachineSerializer
+        elif 'config_context' in request.query_params.get('exclude', []):
+            return serializers.VirtualMachineSerializer
+
+        return serializers.VirtualMachineWithConfigContextSerializer
 
 
 class InterfaceViewSet(ModelViewSet):

netbox/virtualization/tests/test_api.py

@@ -337,6 +337,14 @@ def setUp(self):
         self.virtualmachine1 = VirtualMachine.objects.create(name='Test Virtual Machine 1', cluster=self.cluster1)
         self.virtualmachine2 = VirtualMachine.objects.create(name='Test Virtual Machine 2', cluster=self.cluster1)
         self.virtualmachine3 = VirtualMachine.objects.create(name='Test Virtual Machine 3', cluster=self.cluster1)
+        self.virtualmachine_with_context_data = VirtualMachine.objects.create(
+            name='VM with context data',
+            cluster=self.cluster1,
+            local_context_data={
+                'A': 1,
+                'B': 2
+            }
+        )
 
     def test_get_virtualmachine(self):
 
@@ -350,7 +358,7 @@ def test_list_virtualmachines(self):
         url = reverse('virtualization-api:virtualmachine-list')
         response = self.client.get(url, **self.header)
 
-        self.assertEqual(response.data['count'], 3)
+        self.assertEqual(response.data['count'], 4)
 
     def test_list_virtualmachines_brief(self):
 
@@ -373,7 +381,7 @@ def test_create_virtualmachine(self):
         response = self.client.post(url, data, format='json', **self.header)
 
         self.assertHttpStatus(response, status.HTTP_201_CREATED)
-        self.assertEqual(VirtualMachine.objects.count(), 4)
+        self.assertEqual(VirtualMachine.objects.count(), 5)
         virtualmachine4 = VirtualMachine.objects.get(pk=response.data['id'])
         self.assertEqual(virtualmachine4.name, data['name'])
         self.assertEqual(virtualmachine4.cluster.pk, data['cluster'])
@@ -388,7 +396,7 @@ def test_create_virtualmachine_without_cluster(self):
         response = self.client.post(url, data, format='json', **self.header)
 
         self.assertHttpStatus(response, status.HTTP_400_BAD_REQUEST)
-        self.assertEqual(VirtualMachine.objects.count(), 3)
+        self.assertEqual(VirtualMachine.objects.count(), 4)
 
     def test_create_virtualmachine_bulk(self):
 
@@ -411,7 +419,7 @@ def test_create_virtualmachine_bulk(self):
         response = self.client.post(url, data, format='json', **self.header)
 
         self.assertHttpStatus(response, status.HTTP_201_CREATED)
-        self.assertEqual(VirtualMachine.objects.count(), 6)
+        self.assertEqual(VirtualMachine.objects.count(), 7)
         self.assertEqual(response.data[0]['name'], data[0]['name'])
         self.assertEqual(response.data[1]['name'], data[1]['name'])
         self.assertEqual(response.data[2]['name'], data[2]['name'])
@@ -438,7 +446,7 @@ def test_update_virtualmachine(self):
         response = self.client.put(url, data, format='json', **self.header)
 
         self.assertHttpStatus(response, status.HTTP_200_OK)
-        self.assertEqual(VirtualMachine.objects.count(), 3)
+        self.assertEqual(VirtualMachine.objects.count(), 4)
         virtualmachine1 = VirtualMachine.objects.get(pk=response.data['id'])
         self.assertEqual(virtualmachine1.name, data['name'])
         self.assertEqual(virtualmachine1.cluster.pk, data['cluster'])
@@ -451,7 +459,22 @@ def test_delete_virtualmachine(self):
         response = self.client.delete(url, **self.header)
 
         self.assertHttpStatus(response, status.HTTP_204_NO_CONTENT)
-        self.assertEqual(VirtualMachine.objects.count(), 2)
+        self.assertEqual(VirtualMachine.objects.count(), 3)
+
+    def test_config_context_included_by_default_in_list_view(self):
+
+        url = reverse('virtualization-api:virtualmachine-list')
+        url = '{}?id__in={}'.format(url, self.virtualmachine_with_context_data.pk)
+        response = self.client.get(url, **self.header)
+
+        self.assertEqual(response.data['results'][0].get('config_context', {}).get('A'), 1)
+
+    def test_config_context_excluded(self):
+
+        url = reverse('virtualization-api:virtualmachine-list') + '?exclude=config_context'
+        response = self.client.get(url, **self.header)
+
+        self.assertFalse('config_context' in response.data['results'][0])
 
 
 class InterfaceTest(APITestCase):