Skip to content

Commit 31e6564

Browse files
DragnEmperorDragnEmperor
committed
[fixes] Fixed issue of attached location updating with last_ip
Signed-off-by: DragnEmperor <[email protected]>
1 parent 2c446de commit 31e6564

File tree

2 files changed

+153
-46
lines changed

2 files changed

+153
-46
lines changed

‎openwisp_controller/geo/estimated_location/tasks.py

Lines changed: 82 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -25,40 +25,76 @@ def manage_estimated_locations(device_pk, ip_address, add_existing=False):
2525
- A new location is created if no location exists for current device, or
2626
existing one is updated using coords from WHOIS record if it is estimated (fuzzy).
2727
"""
28+
Device = load_model("config", "Device")
29+
Location = load_model("geo", "Location")
30+
WHOISInfo = load_model("config", "WHOISInfo")
31+
DeviceLocation = load_model("geo", "DeviceLocation")
32+
33+
def _create_estimated_location(device_location, location_defaults):
34+
with transaction.atomic():
35+
location = Location(**location_defaults, is_estimated=True)
36+
location.full_clean()
37+
location.save()
38+
device_location.location = location
39+
device_location.full_clean()
40+
device_location.save()
41+
logger.info(
42+
f"Estimated location saved successfully for {device_pk}"
43+
f" for IP: {ip_address}"
44+
)
2845

29-
def _update_or_create_estimated_location(device_location, whois_obj):
46+
def _update_or_create_estimated_location(
47+
device_location, whois_obj, attached_devices_exists=False
48+
):
3049
# Used to update an existing location if it is estimated
3150
# or create a new one if it doesn't exist
32-
location_defaults = {
33-
**whois_obj._get_defaults_for_estimated_location(),
34-
"organization_id": device.organization_id,
35-
}
36-
if current_location and current_location.is_estimated:
37-
for attr, value in location_defaults.items():
38-
setattr(current_location, attr, value)
39-
current_location.full_clean()
40-
current_location.save()
41-
elif not current_location:
42-
with transaction.atomic():
43-
location = Location(**location_defaults, is_estimated=True)
44-
location.full_clean()
45-
location.save()
46-
device_location.location = location
47-
device_location.full_clean()
48-
device_location.save()
51+
if whois_obj and whois_obj.coordinates:
52+
location_defaults = {
53+
**whois_obj._get_defaults_for_estimated_location(),
54+
"organization_id": device.organization_id,
55+
}
56+
if current_location and current_location.is_estimated:
57+
if attached_devices_exists:
58+
# If there are other devices attached to the current location,
59+
# we do not update it, but create a new one.
60+
_create_estimated_location(device_location, location_defaults)
61+
return
62+
updated = False
63+
for attr, value in location_defaults.items():
64+
if getattr(current_location, attr) != value:
65+
setattr(current_location, attr, value)
66+
updated = True
67+
if updated:
68+
current_location.full_clean()
69+
current_location.save()
70+
logger.info(
71+
f"Estimated location saved successfully for {device_pk}"
72+
f" for IP: {ip_address}"
73+
)
74+
elif not current_location:
75+
# If there is no current location, we create a new one.
76+
_create_estimated_location(device_location, location_defaults)
77+
else:
78+
logger.warning(
79+
f"Coordinates not available for {device_pk} for IP: {ip_address}."
80+
" Estimated location cannot be determined."
81+
)
82+
return
4983

50-
def _handle_attach_existing_location(device, device_location, whois_obj):
84+
def _handle_attach_existing_location(
85+
device, device_location, whois_obj, attached_devices_exists=False
86+
):
5187
# For handling the case when WHOIS already exists for device's new last_ip
5288
# then we attach the location of the device with same last_ip if it exists.
53-
existing_devices_location = (
89+
devices_with_location = (
5490
Device.objects.select_related("devicelocation")
5591
.filter(organization_id=device.organization_id)
5692
.filter(last_ip=ip_address, devicelocation__location__isnull=False)
5793
.exclude(pk=device_pk)
5894
)
5995
# If there are multiple devices with same last_ip then we need to inform
6096
# the user to resolve the conflict manually.
61-
if existing_devices_location.count() > 1:
97+
if devices_with_location.count() > 1:
6298
send_whois_task_notification(
6399
device_pk=device_pk, notify_type="location_error"
64100
)
@@ -67,44 +103,46 @@ def _handle_attach_existing_location(device, device_location, whois_obj):
67103
f"last_ip {ip_address}. Please resolve the conflict manually."
68104
)
69105
return
106+
first_device = devices_with_location.first()
70107
# If existing devices with same last_ip do not have any location
71108
# then we create a new location based on WHOIS data.
72-
if existing_devices_location.count() == 0:
73-
_update_or_create_estimated_location(device_location, whois_obj)
109+
if not first_device:
110+
_update_or_create_estimated_location(
111+
device_location, whois_obj, attached_devices_exists
112+
)
74113
return
75-
existing_location = existing_devices_location.first().devicelocation.location
114+
existing_location = first_device.devicelocation.location
76115
# We need to remove any existing estimated location of the device
77-
if current_location and current_location.pk != existing_location.pk:
116+
if current_location and not attached_devices_exists:
78117
current_location.delete()
79118
device_location.location = existing_location
80119
device_location.full_clean()
81120
device_location.save()
121+
logger.info(
122+
f"Estimated location saved successfully for {device_pk}"
123+
f" for IP: {ip_address}"
124+
)
82125

83-
Device = load_model("config", "Device")
84-
Location = load_model("geo", "Location")
85-
WHOISInfo = load_model("config", "WHOISInfo")
86-
DeviceLocation = load_model("geo", "DeviceLocation")
87-
88-
device = Device.objects.get(pk=device_pk)
89126
whois_obj = WHOISInfo.objects.filter(ip_address=ip_address).first()
127+
device = Device.objects.get(pk=device_pk)
90128
device_location, _ = DeviceLocation.objects.select_related(
91129
"location"
92130
).get_or_create(content_object_id=device_pk)
93-
current_location = device_location.location
94131

95-
if add_existing and (not current_location or current_location.is_estimated):
96-
_handle_attach_existing_location(device, device_location, whois_obj)
132+
attached_devices_exists = False
133+
if current_location := device_location.location:
134+
attached_devices_exists = (
135+
Device.objects.filter(devicelocation__location_id=current_location.pk)
136+
.exclude(pk=device_pk)
137+
.exists()
138+
)
97139

98-
elif whois_obj and whois_obj.coordinates:
99-
_update_or_create_estimated_location(device_location, whois_obj)
140+
if add_existing and (not current_location or current_location.is_estimated):
141+
_handle_attach_existing_location(
142+
device, device_location, whois_obj, attached_devices_exists
143+
)
100144

101145
else:
102-
logger.warning(
103-
f"Coordinates not available for {device_pk} for IP: {ip_address}."
104-
" Estimated location cannot be determined."
146+
_update_or_create_estimated_location(
147+
device_location, whois_obj, attached_devices_exists
105148
)
106-
return
107-
108-
logger.info(
109-
f"Estimated location saved successfully for {device_pk} for IP: {ip_address}"
110-
)

‎openwisp_controller/geo/estimated_location/test_estimated_location.py

Lines changed: 71 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,12 @@ class TestEstimatedLocationTransaction(
103103
_WHOIS_GEOIP_CLIENT = (
104104
"openwisp_controller.config.whois.tasks.geoip2_webservice.Client"
105105
)
106+
_ESTIMATED_LOCATION_INFO_LOGGER = (
107+
"openwisp_controller.geo.estimated_location.tasks.logger.info"
108+
)
109+
_ESTIMATED_LOCATION_ERROR_LOGGER = (
110+
"openwisp_controller.geo.estimated_location.tasks.logger.error"
111+
)
106112

107113
def setUp(self):
108114
super().setUp()
@@ -138,8 +144,9 @@ def test_estimated_location_task_called(
138144
mocked_estimated_location_task.reset_mock()
139145

140146
@mock.patch.object(config_app_settings, "WHOIS_CONFIGURED", True)
147+
@mock.patch(_ESTIMATED_LOCATION_INFO_LOGGER)
141148
@mock.patch(_WHOIS_GEOIP_CLIENT)
142-
def test_estimated_location_creation_and_update(self, mock_client):
149+
def test_estimated_location_creation_and_update(self, mock_client, mock_info):
143150
connect_whois_handlers()
144151

145152
def _verify_location_details(device, mocked_response):
@@ -180,6 +187,11 @@ def _verify_location_details(device, mocked_response):
180187
self.assertEqual(location.is_mobile, False)
181188
self.assertEqual(location.type, "outdoor")
182189
_verify_location_details(device, mocked_response)
190+
mock_info.assert_called_once_with(
191+
f"Estimated location saved successfully for {device.pk}"
192+
f" for IP: {device.last_ip}"
193+
)
194+
mock_info.reset_mock()
183195

184196
with self.subTest("Test Estimated location updated when last ip is updated"):
185197
device.last_ip = "172.217.22.10"
@@ -196,6 +208,11 @@ def _verify_location_details(device, mocked_response):
196208
self.assertEqual(location.is_mobile, False)
197209
self.assertEqual(location.type, "outdoor")
198210
_verify_location_details(device, mocked_response)
211+
mock_info.assert_called_once_with(
212+
f"Estimated location saved successfully for {device.pk}"
213+
f" for IP: {device.last_ip}"
214+
)
215+
mock_info.reset_mock()
199216

200217
with self.subTest(
201218
"Test Location not updated if it is not estimated when last ip is updated"
@@ -213,12 +230,15 @@ def _verify_location_details(device, mocked_response):
213230
self.assertEqual(location.is_mobile, False)
214231
self.assertEqual(location.type, "outdoor")
215232
_verify_location_details(device, mocked_response)
233+
mock_info.assert_not_called()
234+
mock_info.reset_mock()
216235

217236
with self.subTest(
218237
"Test location shared for same IP when new device's location does not exist"
219238
):
220239
Device.objects.all().delete()
221240
device1 = self._create_device(last_ip="172.217.22.10")
241+
mock_info.reset_mock()
222242
device2 = self._create_device(
223243
name="11:22:33:44:55:66",
224244
mac_address="11:22:33:44:55:66",
@@ -228,6 +248,11 @@ def _verify_location_details(device, mocked_response):
228248
self.assertEqual(
229249
device1.devicelocation.location.pk, device2.devicelocation.location.pk
230250
)
251+
mock_info.assert_called_once_with(
252+
f"Estimated location saved successfully for {device2.pk}"
253+
f" for IP: {device2.last_ip}"
254+
)
255+
mock_info.reset_mock()
231256

232257
with self.subTest(
233258
"Test location shared for same IP when new device's location is estimated"
@@ -239,15 +264,21 @@ def _verify_location_details(device, mocked_response):
239264
mac_address="11:22:33:44:55:66",
240265
last_ip="172.217.22.11",
241266
)
267+
mock_info.reset_mock()
242268
old_location = device2.devicelocation.location
243269
device2.last_ip = "172.217.22.10"
244270
device2.save()
271+
mock_info.assert_called_once_with(
272+
f"Estimated location saved successfully for {device2.pk}"
273+
f" for IP: {device2.last_ip}"
274+
)
245275
device2.refresh_from_db()
246276

247277
self.assertEqual(
248278
device1.devicelocation.location.pk, device2.devicelocation.location.pk
249279
)
250280
self.assertEqual(Location.objects.filter(pk=old_location.pk).count(), 0)
281+
mock_info.reset_mock()
251282

252283
with self.subTest(
253284
"Test location not shared for same IP when new "
@@ -260,21 +291,53 @@ def _verify_location_details(device, mocked_response):
260291
mac_address="11:22:33:44:55:66",
261292
last_ip="172.217.22.11",
262293
)
294+
mock_info.reset_mock()
263295
old_location = device2.devicelocation.location
264296
old_location.is_estimated = False
265297
old_location.save()
266298
device2.last_ip = "172.217.22.10"
267299
device2.save()
300+
mock_info.assert_not_called()
268301
device2.refresh_from_db()
269302

270303
self.assertNotEqual(
271304
device1.devicelocation.location.pk, device2.devicelocation.location.pk
272305
)
273306
self.assertEqual(Location.objects.filter(pk=old_location.pk).count(), 1)
307+
mock_info.reset_mock()
308+
309+
with self.subTest(
310+
"Shared location not updated when either device's last_ip changes. "
311+
"New location created for device with updated last_ip"
312+
):
313+
Device.objects.all().delete()
314+
device1 = self._create_device(last_ip="172.217.22.10")
315+
device2 = self._create_device(
316+
name="11:22:33:44:55:66",
317+
mac_address="11:22:33:44:55:66",
318+
last_ip="172.217.22.10",
319+
)
320+
mock_info.reset_mock()
321+
self.assertEqual(
322+
device1.devicelocation.location.pk, device2.devicelocation.location.pk
323+
)
324+
device2.last_ip = "172.217.22.11"
325+
device2.save()
326+
mock_info.assert_called_once_with(
327+
f"Estimated location saved successfully for {device2.pk}"
328+
f" for IP: {device2.last_ip}"
329+
)
330+
device2.refresh_from_db()
331+
self.assertNotEqual(
332+
device1.devicelocation.location.pk, device2.devicelocation.location.pk
333+
)
334+
mock_info.reset_mock()
274335

275336
@mock.patch.object(config_app_settings, "WHOIS_CONFIGURED", True)
337+
@mock.patch(_ESTIMATED_LOCATION_INFO_LOGGER)
338+
@mock.patch(_ESTIMATED_LOCATION_ERROR_LOGGER)
276339
@mock.patch(_WHOIS_GEOIP_CLIENT)
277-
def test_estimated_location_notification(self, mock_client):
340+
def test_estimated_location_notification(self, mock_client, mock_error, mock_info):
278341
"""
279342
For testing notification related to location is sent to user
280343
when already multiple devices with same last_ip exist.
@@ -288,13 +351,19 @@ def test_estimated_location_notification(self, mock_client):
288351
mac_address="11:22:33:44:55:66",
289352
last_ip="172.217.22.10",
290353
)
354+
mock_info.reset_mock()
291355
# device3 will not have same location as first two devices
292356
# as multiple devices found for same last_ip causing conflict
293357
device3 = self._create_device(
294358
name="11:22:33:44:55:77",
295359
mac_address="11:22:33:44:55:77",
296360
last_ip="172.217.22.10",
297361
)
362+
mock_info.assert_not_called()
363+
mock_error.assert_called_once_with(
364+
f"Multiple devices with locations found with same "
365+
f"last_ip {device3.last_ip}. Please resolve the conflict manually."
366+
)
298367
self.assertEqual(notification_qs.count(), 1)
299368
notification = notification_qs.first()
300369
self.assertEqual(notification.actor, device3)

0 commit comments

Comments
 (0)