[changes] Assert queries count in estimated location tests

Signed-off-by: DragnEmperor <[email protected]>

Author: DragnEmperor

openwisp_controller/config/whois/service.py

@@ -305,7 +305,7 @@ def _create_or_update_estimated_location(
                 device_location.full_clean()
                 device_location.save()
                 send_whois_task_notification(
-                    device_pk=self.device.pk,
+                    device=self.device,
                     notify_type="estimated_location_created",
                     actor=current_location,
                 )
@@ -322,7 +322,7 @@ def _create_or_update_estimated_location(
                     )
 
                 send_whois_task_notification(
-                    device_pk=self.device.pk,
+                    device=self.device,
                     notify_type="estimated_location_updated",
                     actor=current_location,
                 )

openwisp_controller/config/whois/tasks.py

@@ -26,9 +26,7 @@ class WHOISCeleryRetryTask(OpenwispCeleryTask):
     def on_failure(self, exc, task_id, args, kwargs, einfo):
         """Notify the user about the failure of the WHOIS task."""
         device_pk = kwargs.get("device_pk")
-        send_whois_task_notification(
-            device_pk=device_pk, notify_type="whois_device_error"
-        )
+        send_whois_task_notification(device=device_pk, notify_type="whois_device_error")
         logger.error(f"WHOIS lookup failed. Details: {exc}")
         return super().on_failure(exc, task_id, args, kwargs, einfo)
 
@@ -51,7 +49,7 @@ def _manage_whois_record(whois_details, whois_instance=None):
     else:
         whois_instance = WHOISInfo(**whois_details)
         whois_instance.full_clean()
-        whois_instance.save()
+        whois_instance.save(force_insert=True)
     return whois_instance, update_fields
 
 

openwisp_controller/config/whois/tests/utils.py

@@ -61,10 +61,14 @@ def _task_called(self, mocked_task, task_name="WHOIS lookup"):
         org = self._get_org()
 
         with self.subTest(f"{task_name} task called when last_ip is public"):
-            with mock.patch("django.core.cache.cache.set") as mocked_set:
+            with mock.patch("django.core.cache.cache.get") as mocked_get, mock.patch(
+                "django.core.cache.cache.set"
+            ) as mocked_set:
+                mocked_get.side_effect = [None, org.config_settings]
                 device = self._create_device(last_ip="172.217.22.14")
                 mocked_task.assert_called()
                 mocked_set.assert_called_once()
+                mocked_get.assert_called()
         mocked_task.reset_mock()
 
         with self.subTest(
@@ -77,7 +81,7 @@ def _task_called(self, mocked_task, task_name="WHOIS lookup"):
                 device.save()
                 mocked_task.assert_called()
                 mocked_set.assert_not_called()
-                mocked_get.assert_called_once()
+                mocked_get.assert_called()
         mocked_task.reset_mock()
 
         with self.subTest(f"{task_name} task not called when last_ip not updated"):

openwisp_controller/config/whois/utils.py

@@ -41,10 +41,10 @@
 }
 
 
-def send_whois_task_notification(device_pk, notify_type, actor=None):
+def send_whois_task_notification(device, notify_type, actor=None):
     Device = load_model("config", "Device")
-
-    device = Device.objects.get(pk=device_pk)
+    if not isinstance(device, Device):
+        device = Device.objects.get(pk=device)
     notify_details = MESSAGE_MAP[notify_type]
     notify.send(
         sender=actor or device,

openwisp_controller/geo/estimated_location/tasks.py

@@ -32,17 +32,17 @@ def _handle_attach_existing_location(
         existing_location = existing_device_location.location
         # We need to remove existing estimated location of the device
         # if it is not shared
-        if attached_devices_exists is False:
-            current_location.delete()
         device_location.location = existing_location
         device_location.full_clean()
         device_location.save()
         logger.info(
             f"Estimated location saved successfully for {device.pk}"
             f" for IP: {ip_address}"
         )
+        if attached_devices_exists is False:
+            current_location.delete()
         send_whois_task_notification(
-            device_pk=device.pk,
+            device=device,
             notify_type="estimated_location_updated",
             actor=existing_location,
         )
@@ -94,23 +94,19 @@ def manage_estimated_locations(device_pk, ip_address):
     Device = load_model("config", "Device")
     DeviceLocation = load_model("geo", "DeviceLocation")
 
-    device = (
-        Device.objects.select_related("devicelocation__location", "organization")
-        .only("organization_id", "devicelocation")
-        .get(pk=device_pk)
-    )
+    device = Device.objects.select_related("devicelocation__location").get(pk=device_pk)
 
     devices_with_location = (
-        Device.objects.select_related("devicelocation")
+        Device.objects.only("devicelocation")
+        .select_related("devicelocation__location")
         .filter(organization_id=device.organization_id)
         .filter(last_ip=ip_address, devicelocation__location__isnull=False)
         .exclude(pk=device.pk)
     )
-    # If there are multiple devices with same last_ip then we need to inform
-    # the user to resolve the conflict manually.
+
     if devices_with_location.count() > 1:
         send_whois_task_notification(
-            device_pk=device.pk, notify_type="estimated_location_error"
+            device=device, notify_type="estimated_location_error"
         )
         logger.error(
             "Multiple devices with locations found with same "

openwisp_controller/geo/estimated_location/tests/tests.py

@@ -18,6 +18,7 @@
 from ....tests.utils import TestAdminMixin
 from ...tests.utils import TestGeoMixin
 from ..handlers import register_estimated_location_notification_types
+from ..tasks import manage_estimated_locations
 from .utils import TestEstimatedLocationMixin
 
 Device = load_model("config", "Device")
@@ -282,9 +283,20 @@ def test_estimated_location_task_called(
         mocked_estimated_location_task.reset_mock()
 
     @mock.patch.object(config_app_settings, "WHOIS_CONFIGURED", True)
+    @mock.patch(
+        "openwisp_controller.config.whois.service.send_whois_task_notification"  # noqa
+    )
+    @mock.patch(
+        "openwisp_controller.geo.estimated_location.tasks.send_whois_task_notification"  # noqa
+    )
+    @mock.patch(
+        "openwisp_controller.geo.estimated_location.tasks.manage_estimated_locations.delay"  # noqa
+    )
     @mock.patch(_ESTIMATED_LOCATION_INFO_LOGGER)
     @mock.patch(_WHOIS_GEOIP_CLIENT)
-    def test_estimated_location_creation_and_update(self, mock_client, mock_info):
+    def test_estimated_location_creation_and_update(
+        self, mock_client, mock_info, _mocked_task, _mocked_notify, _mocked_notify2
+    ):
         connect_whois_handlers()
 
         def _verify_location_details(device, mocked_response):
@@ -318,6 +330,8 @@ def _verify_location_details(device, mocked_response):
 
         with self.subTest("Test Estimated location created when device is created"):
             device = self._create_device(last_ip="172.217.22.14")
+            with self.assertNumQueries(14):
+                manage_estimated_locations(device.pk, device.last_ip)
 
             location = device.devicelocation.location
             mocked_response.ip_address = device.last_ip
@@ -338,6 +352,8 @@ def _verify_location_details(device, mocked_response):
             mocked_response.city.name = "New City"
             mock_client.return_value.city.return_value = mocked_response
             device.save()
+            with self.assertNumQueries(8):
+                manage_estimated_locations(device.pk, device.last_ip)
             device.refresh_from_db()
 
             location = device.devicelocation.location
@@ -361,6 +377,8 @@ def _verify_location_details(device, mocked_response):
             mock_client.return_value.city.return_value = self._mocked_client_response()
             device.devicelocation.location.save(_set_estimated=True)
             device.save()
+            with self.assertNumQueries(2):
+                manage_estimated_locations(device.pk, device.last_ip)
             device.refresh_from_db()
 
             location = device.devicelocation.location
@@ -379,12 +397,15 @@ def _verify_location_details(device, mocked_response):
         ):
             Device.objects.all().delete()
             device1 = self._create_device(last_ip="172.217.22.10")
+            manage_estimated_locations(device1.pk, device1.last_ip)
             mock_info.reset_mock()
             device2 = self._create_device(
                 name="11:22:33:44:55:66",
                 mac_address="11:22:33:44:55:66",
                 last_ip="172.217.22.10",
             )
+            with self.assertNumQueries(8):
+                manage_estimated_locations(device2.pk, device2.last_ip)
 
             self.assertEqual(
                 device1.devicelocation.location.pk, device2.devicelocation.location.pk
@@ -400,15 +421,20 @@ def _verify_location_details(device, mocked_response):
         ):
             Device.objects.all().delete()
             device1 = self._create_device(last_ip="172.217.22.10")
+            manage_estimated_locations(device1.pk, device1.last_ip)
             device2 = self._create_device(
                 name="11:22:33:44:55:66",
                 mac_address="11:22:33:44:55:66",
                 last_ip="172.217.22.11",
             )
+            manage_estimated_locations(device2.pk, device2.last_ip)
             mock_info.reset_mock()
             old_location = device2.devicelocation.location
             device2.last_ip = "172.217.22.10"
             device2.save()
+            # 3 queries related to notifications cleanup
+            with self.assertNumQueries(16):
+                manage_estimated_locations(device2.pk, device2.last_ip)
             mock_info.assert_called_once_with(
                 f"Estimated location saved successfully for {device2.pk}"
                 f" for IP: {device2.last_ip}"
@@ -427,17 +453,21 @@ def _verify_location_details(device, mocked_response):
         ):
             Device.objects.all().delete()
             device1 = self._create_device(last_ip="172.217.22.10")
+            manage_estimated_locations(device1.pk, device1.last_ip)
             device2 = self._create_device(
                 name="11:22:33:44:55:66",
                 mac_address="11:22:33:44:55:66",
                 last_ip="172.217.22.11",
             )
+            manage_estimated_locations(device2.pk, device2.last_ip)
             mock_info.reset_mock()
             old_location = device2.devicelocation.location
             old_location.is_estimated = False
             old_location.save()
             device2.last_ip = "172.217.22.10"
             device2.save()
+            with self.assertNumQueries(2):
+                manage_estimated_locations(device2.pk, device2.last_ip)
             mock_info.assert_called_once_with(
                 f"Non Estimated location already set for {device2.pk}. Update"
                 f" location manually as per IP: {device2.last_ip}"
@@ -456,17 +486,21 @@ def _verify_location_details(device, mocked_response):
         ):
             Device.objects.all().delete()
             device1 = self._create_device(last_ip="172.217.22.10")
+            manage_estimated_locations(device1.pk, device1.last_ip)
             device2 = self._create_device(
                 name="11:22:33:44:55:66",
                 mac_address="11:22:33:44:55:66",
                 last_ip="172.217.22.10",
             )
+            manage_estimated_locations(device2.pk, device2.last_ip)
             mock_info.reset_mock()
             self.assertEqual(
                 device1.devicelocation.location.pk, device2.devicelocation.location.pk
             )
             device2.last_ip = "172.217.22.11"
             device2.save()
+            with self.assertNumQueries(14):
+                manage_estimated_locations(device2.pk, device2.last_ip)
             mock_info.assert_called_once_with(
                 f"Estimated location saved successfully for {device2.pk}"
                 f" for IP: {device2.last_ip}"