Fixing tests.

WoTTsecurity · Aug 31, 2019 · 1973b7d · 1973b7d
1 parent 781f32d
commit 1973b7d
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/backend/device_registry/migrations/0054_device_calculated_trust_score.py b/backend/device_registry/migrations/0054_device_calculated_trust_score.py
@@ -1,10 +1,10 @@
 # Generated by Django 2.1.10 on 2019-08-06 05:10
 
 from django.db import migrations, models
-from device_registry.models import Device
 
 
 def save_trust_score(apps, schema_editor):
+    Device = apps.get_model('device_registry', 'Device')
     for d in Device.objects.all():
         d.save(update_fields=['trust_score'])
 

diff --git a/backend/device_registry/tests/test_api.py b/backend/device_registry/tests/test_api.py
@@ -717,15 +717,26 @@ def setUp(self):
     def test_ping_get_success(self):
         response = self.client.get(self.url, **self.headers)
         self.assertEqual(response.status_code, 200)
-        self.assertDictEqual(response.data, {'policy': self.device.firewallstate.policy_string,
-                                             'block_ports': [], 'block_networks': settings.SPAM_NETWORKS})
+        self.assertDictEqual(response.data, {
+            'firewall': {
+                'policy': self.device.firewallstate.policy_string,
+                'block_ports': [], 'block_networks': settings.SPAM_NETWORKS
+            },
+            'deb_packages_hash': None
+        })
 
     def test_pong_data(self):
         # 1st request
         response = self.client.get(self.url, **self.headers)
         self.assertEqual(response.status_code, 200)
-        self.assertDictEqual(response.data, {'block_ports': [], 'block_networks': settings.SPAM_NETWORKS,
-                                             'policy': self.device.firewallstate.policy_string})
+        self.assertDictEqual(response.data, {
+            'firewall': {
+                'block_ports': [], 
+                'block_networks': settings.SPAM_NETWORKS,
+                'policy': self.device.firewallstate.policy_string
+            },
+            'deb_packages_hash': None
+        })
         # 2nd request
         self.device.portscan.block_ports = [['192.168.1.178', 'tcp', 22, False]]
         self.device.portscan.block_networks = [['192.168.1.177', False]]
@@ -735,9 +746,14 @@ def test_pong_data(self):
 
         response = self.client.get(self.url, **self.headers)
         self.assertEqual(response.status_code, 200)
-        self.assertDictEqual(response.data, {'policy': self.device.firewallstate.policy_string,
-                                             'block_ports': [['192.168.1.178', 'tcp', 22, False]],
-                                             'block_networks': [['192.168.1.177', False]] + settings.SPAM_NETWORKS})
+        self.assertDictEqual(response.data, {
+            'firewall': {
+                'policy': self.device.firewallstate.policy_string,
+                'block_ports': [['192.168.1.178', 'tcp', 22, False]],
+                'block_networks': [['192.168.1.177', False]] + settings.SPAM_NETWORKS
+            },
+            'deb_packages_hash': None
+        })
 
     def test_ping_creates_models(self):
         devinfo_obj_count_before = DeviceInfo.objects.count()