Do not patch or delete device operations on ambiguous device uuids

Change-type: patch
balena-io · Oct 15, 2024 · e82b65e · e82b65e
1 parent d8fbe7e
commit e82b65e
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 1 deletion.
diff --git a/balena/models/device.py b/balena/models/device.py
@@ -174,6 +174,19 @@ def __set(
                     }
                 )
             else:
+                if len(uuid_or_id_or_ids) < 6:
+                    raise exceptions.AmbiguousDevice(uuid_or_id_or_ids)
+
+                affected = self.__pine.get({
+                    "resource": "device",
+                    "options": {
+                        "$count": {
+                           "$filter": {"uuid": {"$startswith": uuid_or_id_or_ids}}}
+                        }
+                    })
+                if affected > 1:
+                    raise exceptions.AmbiguousDevice(uuid_or_id_or_ids)
+
                 fn(
                     {
                         "resource": "device",
@@ -576,7 +589,6 @@ def remove(self, uuid_or_id_or_ids: Union[str, int, List[int]]):
         Args:
             uuid_or_id_or_ids (Union[str, int, List[int]]): device uuid (str) or id (int) or ids (List[int])
         """
-
         self.__set(uuid_or_id_or_ids, body=None, fn=self.__pine.delete)
 
     def deactivate(self, uuid_or_id_or_ids: Union[str, int, List[int]]) -> None:

diff --git a/tests/functional/models/test_device.py b/tests/functional/models/test_device.py
@@ -369,9 +369,39 @@ def test_30_remove(self):
             app_devices_len - 1,
         )
 
+        uuid = self.balena.models.device.generate_uuid()
+        self.balena.models.device.register(self.app["id"], uuid)
+        uuid2 = uuid[:-2] + uuid[-1] + uuid[-2]
+
+        self.balena.models.device.register(self.app["id"], uuid2)
+
+        device_uuids = [device["uuid"] for device in self.balena.models.device.get_all()]
+        self.assertIn(uuid, device_uuids)
+        self.assertIn(uuid2, device_uuids)
+
         with self.assertRaises(self.helper.balena_exceptions.DeviceNotFound):
             self.balena.models.device.get(TestDevice.device["uuid"])
 
+        with self.assertRaises(self.helper.balena_exceptions.AmbiguousDevice):
+            self.balena.models.device.remove("")
+
+        with self.assertRaises(self.helper.balena_exceptions.AmbiguousDevice):
+            self.balena.models.device.remove("abc")
+
+        with self.assertRaises(self.helper.balena_exceptions.AmbiguousDevice):
+            self.balena.models.device.remove(uuid[0:10])
+
+        device_uuids = [device["uuid"] for device in self.balena.models.device.get_all()]
+        self.assertIn(uuid, device_uuids)
+        self.assertIn(uuid2, device_uuids)
+
+        self.balena.models.device.remove(uuid)
+        self.balena.models.device.remove(uuid2)
+
+        device_uuids = [device["uuid"] for device in self.balena.models.device.get_all()]
+        self.assertNotIn(uuid, device_uuids)
+        self.assertNotIn(uuid2, device_uuids)
+
 
 if __name__ == "__main__":
     unittest.main()