better organise code

device/tests/test_mock_device.py

@@ -3,30 +3,14 @@ from unittest.mock import MagicMock
 
 
 class TestDevice(Device):
-    """A test subclass of Device that overrides the database-dependent methods"""
+    def __init__(self, id):
-    # TODO Leaving commented bc not used, but might be useful at some  point
+        super().__init__(id=id)
-    # def get_annotations(self):
+        self.shortid = id[:6].upper()
-    #     """Return empty list instead of querying database"""
+        self.uuids = []
-    #     return []
+        self.hids = ['hid1', 'hid2']
+        self._setup_evidence()
 
-    # def get_uuids(self):
+    def _setup_evidence(self):
-    #     """Set uuids directly instead of querying"""
-    #     self.uuids = ['uuid1', 'uuid2']
-
-    # def get_hids(self):
-    #     """Set hids directly instead of querying"""
-    #     self.hids = ['hid1', 'hid2']
-
-    # def get_evidences(self):
-    #     """Set evidences directly instead of querying"""
-    #     self.evidences = []
-
-    # def get_lots(self):
-    #     """Set lots directly instead of querying"""
-    #     self.lots = []
-
-    def get_last_evidence(self):
-        if not hasattr(self, '_evidence'):
         self._evidence = MagicMock()
         self._evidence.doc = {
             'type': 'Computer',
@@ -54,23 +38,10 @@ class TestDevice(Device):
         ]
         self.last_evidence = self._evidence
 
+    @property
+    def components(self):
+        return self.last_evidence.get_components()
 
-class TestWebSnapshotDevice(TestDevice):
+    @property
-    """A test subclass of Device that simulates a WebSnapshot device"""
+    def serial_number(self):
-
+        return self.last_evidence.doc['device']['serialNumber']
-    def get_last_evidence(self):
-        if not hasattr(self, '_evidence'):
-            self._evidence = MagicMock()
-            self._evidence.doc = {
-                'type': 'WebSnapshot',
-                'kv': {
-                    'URL': 'http://example.com',
-                    'Title': 'Test Page',
-                    'Timestamp': '2024-01-01'
-                },
-                'device': {
-                    'type': 'Laptop'
-                }
-            }
-        self.last_evidence = self._evidence
-        return self._evidence

device/tests/test_public_device_web.py

@@ -2,8 +2,8 @@ from django.test import TestCase, Client
 from django.urls import reverse
 from unittest.mock import patch
 from device.views import PublicDeviceWebView
-from device.tests.test_mock_device import TestDevice, TestWebSnapshotDevice
+from device.tests.test_mock_device import TestDevice
-from user.models import User, Institution  # Import both models
+from user.models import User, Institution
 
 
 class PublicDeviceWebViewTests(TestCase):
@@ -22,7 +22,6 @@ class PublicDeviceWebViewTests(TestCase):
         )
 
     def test_url_resolves_correctly(self):
-        """Test that the URL is constructed correctly"""
         url = reverse('device:device_web', kwargs={'pk': self.test_id})
         self.assertEqual(url, f'/device/{self.test_id}/public/')
 
@@ -66,22 +65,13 @@ class PublicDeviceWebViewTests(TestCase):
     @patch('device.views.Device')
     def test_json_response_anonymous(self, MockDevice):
         test_device = TestDevice(id=self.test_id)
-        test_device.get_last_evidence()
+        MockDevice.return_value = test_device
-        mock_instance = MockDevice.return_value
-        mock_instance.id = self.test_id
-        mock_instance.shortid = self.test_id[:6].upper()
-        mock_instance.uuids = []
-        mock_instance.hids = ['hid1', 'hid2']
-        mock_instance.last_evidence = test_device.last_evidence
-
         response = self.client.get(
             self.test_url,
             HTTP_ACCEPT='application/json'
         )
-
         self.assertEqual(response.status_code, 200)
         self.assertEqual(response['Content-Type'], 'application/json')
-
         json_data = response.json()
         self.assertEqual(json_data['id'], self.test_id)
         self.assertEqual(json_data['shortid'], self.test_id[:6].upper())
@@ -93,22 +83,12 @@ class PublicDeviceWebViewTests(TestCase):
     @patch('device.views.Device')
     def test_json_response_authenticated(self, MockDevice):
         test_device = TestDevice(id=self.test_id)
-        test_device.get_last_evidence()
+        MockDevice.return_value = test_device
-        mock_instance = MockDevice.return_value
-        mock_instance.id = self.test_id
-        mock_instance.shortid = self.test_id[:6].upper()
-        mock_instance.uuids = []
-        mock_instance.hids = ['hid1', 'hid2']
-        mock_instance.last_evidence = test_device.last_evidence
-        mock_instance.components = test_device.last_evidence.get_components()
-        mock_instance.serial_number = test_device.last_evidence.doc['device']['serialNumber']
-
         self.client.login(username='test@example.com', password='testpass123')
         response = self.client.get(
             self.test_url,
             HTTP_ACCEPT='application/json'
         )
-
         self.assertEqual(response.status_code, 200)
         self.assertEqual(response['Content-Type'], 'application/json')
 
@@ -130,15 +110,3 @@ class PublicDeviceWebViewTests(TestCase):
         self.assertEqual(json_data['serial_number'], 'SN123456')
         self.assertEqual(json_data['uuids'], [])
         self.assertEqual(json_data['hids'], ['hid1', 'hid2'])
-
-    @patch('device.views.Device')
-    def test_websnapshot_device(self, MockDevice):
-        test_device = TestWebSnapshotDevice(id=self.test_id)
-        MockDevice.return_value = test_device
-        response = self.client.get(self.test_url)
-
-        self.assertEqual(response.status_code, 200)
-        self.assertTemplateUsed(response, 'device_web.html')
-
-        self.assertContains(response, 'http://example.com')
-        self.assertContains(response, 'Test Page')