Refactor Omada integration: encapsulate API calls in Omada class, update device management tasks, and enhance device creation tests

.python-version

@@ -0,0 +1 @@
+3.12.5

api/omada.py

@@ -0,0 +1,93 @@
+import requests
+from apibase.env import env
+
+
+class Omada:
+    def __init__(self):
+        self.proxy_url = env("OMADA_PROXY_URL", default="")  # type: ignore
+        self.api_key = env.str("OMADA_PROXY_API_KEY", default="")  # type: ignore
+        self.site_id = env("OMADA_SITE_ID", default="")  # type: ignore
+        self.group_id = env("OMADA_GROUP_ID", default="")  # type: ignore
+
+        if not self.proxy_url:
+            raise ValueError("OMADA_PROXY_URL is not set in the environment variables.")
+        if not self.api_key:
+            raise ValueError(
+                "OMADA_PROXY_API_KEY is not set in the environment variables."
+            )
+        if not self.site_id:
+            raise ValueError("OMADA_SITE_ID is not set in the environment variables.")
+        if not self.group_id:
+            raise ValueError("OMADA_GROUP_ID is not set in the environment variables.")
+
+    def get_existing_omada_devices(self):
+        """
+        Get existing Omada devices from the database.
+        :return: List of existing device names.
+        """
+        try:
+            url = f"{self.proxy_url}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/{self.site_id}/setting/profiles/groups"
+            response = requests.get(
+                url,
+                headers={"X-API-Key": str(self.api_key)},
+            )
+            print("Response: ", response.status_code)
+            data = response.json()
+            existing_devices = []
+            if "result" in data and len(data["result"]["data"]) > 0:
+                last_entry = data["result"]["data"][-1]
+                print("Last Entry: ", last_entry)
+                if "macAddressList" in last_entry:
+                    existing_devices = last_entry["macAddressList"]
+            print(existing_devices)
+            return existing_devices
+        except requests.RequestException as e:
+            print(f"Error fetching existing devices: {e}")
+            return []
+
+    def add_new_devices_to_omada(self, new_devices: list[dict]):
+        """
+        Add new devices to Omada.
+        :param new_devices: List of new device names to add.
+        """
+        try:
+            PAYLOAD = {
+                "name": "REGISTERED_DEVICES",
+                "type": 2,
+                "resource": 0,
+                "ipList": None,
+                "ipv6List": None,
+                "macAddressList": None,
+                "portList": None,
+                "countryList": None,
+                "portType": None,
+                "portMaskList": None,
+                "domainNamePort": None,
+            }
+            existing_devices = self.get_existing_omada_devices()
+            PAYLOAD["macAddressList"] = existing_devices
+            print("Payload with existing devices: ", PAYLOAD)
+            for device in new_devices:
+                print("Device in loop: ", device)
+                PAYLOAD["macAddressList"].append(
+                    {
+                        "macAddress": device["mac"],
+                        "name": device["name"],
+                    }
+                )
+            print("New Payload: ", PAYLOAD)
+            url = f"{self.proxy_url}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/{self.site_id}/setting/profiles/groups/2/{self.group_id}"
+            print(url)
+            response = requests.patch(
+                url,
+                headers={"X-API-Key": str(self.api_key)},
+                json=PAYLOAD,
+            )
+            print("Response: ", response.status_code)
+            if response.status_code == 200:
+                print("Devices successfully added.")
+                print(response.json())
+            else:
+                print(f"Failed to add devices: {response.text}")
+        except requests.RequestException as e:
+            print(f"Error adding devices to Omada: {e}")

api/tasks.py

@@ -2,23 +2,19 @@ from django.shortcuts import get_object_or_404
 from api.models import User
 from devices.models import Device
 from api.notifications import send_sms
-import requests
-from apibase.env import env, BASE_DIR
 import os
 import logging
 from celery import shared_task
 from django.utils import timezone
 from api.notifications import send_clean_telegram_markdown
+from api.omada import Omada
+from apibase.env import env, BASE_DIR
 
 logger = logging.getLogger(__name__)
 
-
 env.read_env(os.path.join(BASE_DIR, ".env"))
-PERSON_VERIFY_BASE_URL = env.str("PERSON_VERIFY_BASE_URL", default="")  # type: ignore
+
-OMADA_PROXY_API_KEY = env.str("OMADA_PROXY_API_KEY", default="")  # type: ignore
+omada_client = Omada()
-OMADA_PROXY_URL = env("OMADA_PROXY_URL", default="")  # type: ignore
-OMADA_SITE_ID = env("OMADA_SITE_ID", default="")  # type: ignore
-OMADA_GROUP_ID = env("OMADA_GROUP_ID", default="")  # type: ignore
 
 
 @shared_task
@@ -61,93 +57,19 @@ def deactivate_expired_devices():
 
 def get_existing_omada_devices():
     """
-    Get existing Omada devices from the database.
+    Get existing Omada devices from Omada API via Omada class.
     :return: List of existing device names.
     """
-    if not OMADA_PROXY_URL:
+    return omada_client.get_existing_omada_devices()
-        raise ValueError(
-            "OMADA_PROXY_URL is not set. Please set it in your environment variables."
-        )
-    try:
-        response = requests.get(
-            f"{OMADA_PROXY_URL}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/66dcddb804aa0d2978cf145f/setting/profiles/groups",
-            headers={"X-API-Key": str(OMADA_PROXY_API_KEY)},
-        )
-        print("Response: ", response.status_code)
-        data = response.json()
-        existing_devices = []
-        if "result" in data and len(data["result"]["data"]) > 0:
-            last_entry = data["result"]["data"][-1]
-            print("Last Entry: ", last_entry)
-            if "macAddressList" in last_entry:
-                existing_devices = last_entry["macAddressList"]
-        print(existing_devices)
-        return existing_devices
-    except requests.RequestException as e:
-        print(f"Error fetching existing devices: {e}")
-        return []
 
 
 @shared_task
 def add_new_devices_to_omada(new_devices: list[dict]):
     """
-    Add new devices to Omada.
+    Add new devices to Omada via Omada class.
     :param new_devices: List of new device names to add.
     """
-    if not OMADA_SITE_ID:
+    omada_client.add_new_devices_to_omada(new_devices)
-        raise ValueError(
-            "OMADA_SITE_ID is not set. Please set it in your environment variables."
-        )
-    if not OMADA_GROUP_ID:
-        raise ValueError(
-            "OMADA_GROUP_ID is not set. Please set it in your environment variables."
-        )
-    if not OMADA_PROXY_URL:
-        raise ValueError(
-            "OMADA_PROXY_URL is not set. Please set it in your environment variables."
-        )
-    try:
-        PAYLOAD = {
-            "name": "REGISTERED_DEVICES",
-            "type": 2,
-            "resource": 0,
-            "ipList": None,
-            "ipv6List": None,
-            "macAddressList": None,
-            "portList": None,
-            "countryList": None,
-            "portType": None,
-            "portMaskList": None,
-            "domainNamePort": None,
-        }
-        existing_devices = get_existing_omada_devices()
-        PAYLOAD["macAddressList"] = existing_devices
-        print("Payload with existing devices: ", PAYLOAD)
-        for device in new_devices:
-            print("Device in loop: ", device)
-            PAYLOAD["macAddressList"].append(
-                {
-                    "macAddress": device["mac"],
-                    "name": device["name"],
-                }
-            )
-        print("New Payload: ", PAYLOAD)
-        print(
-            f"{OMADA_PROXY_URL}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/{OMADA_SITE_ID}/setting/profiles/groups/2/{OMADA_GROUP_ID}"
-        )
-        response = requests.patch(
-            f"{OMADA_PROXY_URL}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/{OMADA_SITE_ID}/setting/profiles/groups/2/{OMADA_GROUP_ID}",
-            headers={"X-API-Key": str(OMADA_PROXY_API_KEY)},
-            json=PAYLOAD,
-        )
-        print("Response: ", response.status_code)
-        if response.status_code == 200:
-            print("Devices successfully added.")
-            print(response.json())
-        else:
-            print(f"Failed to add devices: {response.text}")
-    except requests.RequestException as e:
-        print(f"Error adding devices: {e}")
 
 
 def verify_user_with_person_api_task(user_id: int):
@@ -155,7 +77,6 @@ def verify_user_with_person_api_task(user_id: int):
     Verify the user with the Person API.
     :param user_id: The ID of the user to verify.
     """
-
     user = get_object_or_404(User, id=user_id)
     # Call the Person API to verify the user
 
@@ -171,10 +92,14 @@ def verify_user_with_person_api_task(user_id: int):
     """
 
     logger.info(verification_failed_message)
+    PERSON_VERIFY_BASE_URL = env.str("PERSON_VERIFY_BASE_URL", default="")  # type: ignore
+
     if not PERSON_VERIFY_BASE_URL:
         raise ValueError(
             "PERSON_VERIFY_BASE_URL is not set in the environment variables."
         )
+    import requests
+
     response = requests.get(f"{PERSON_VERIFY_BASE_URL}/api/person/{user.id_card}")
     if response.status_code == 200:
         data = response.json()
@@ -250,4 +175,4 @@ def verify_user_with_person_api_task(user_id: int):
     else:
         # Handle the error case
         print(f"Error verifying user: {response.status_code} - {response.text}")
-        return False
+        return

devices/tests.py

@@ -35,11 +35,18 @@ class DeviceAPITestCase(APITestCase):
         response = self.client.get("/api/devices/")
         self.assertEqual(response.status_code, status.HTTP_200_OK)
 
-    def test_create_device(self):
+    def test_create_device_with_no_vendor_mac(self):
         self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}")
         data = {"name": "New Device", "mac": "11:22:33:44:55:66"}
         response = self.client.post("/api/devices/", data)
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+    def test_create_device_with_vendor_mac(self):
+        self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}")
+        data = {"name": "New Device", "mac": "64:16:7F:21:51:A5"}
+        response = self.client.post("/api/devices/", data)
         self.assertEqual(response.status_code, status.HTTP_201_CREATED)
+        self.assertIn("New Device", response.data["name"])
 
     def test_create_device_invalid_mac(self):
         self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}")
@@ -48,6 +55,15 @@ class DeviceAPITestCase(APITestCase):
         self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
         self.assertIn("Invalid mac address.", response.data["message"])
 
+    def test_mac_address_already_exists(self):
+        self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}")
+        data = {"name": "Existing Device", "mac": self.device.mac}
+        response = self.client.post("/api/devices/", data)
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+        self.assertIn(
+            "Device with this mac address already exists.", response.data["message"]
+        )
+
     def test_update_device(self):
         self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}")
         data = {"name": "Updated Device"}
@@ -59,14 +75,14 @@ class DeviceAPITestCase(APITestCase):
     def test_block_device(self):
         self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}")
         data = {
-            "blocked": True,  # also use Python boolean here, not a string
+            "blocked": True,
             "reason_for_blocking": "hello",
             "blocked_by": "ADMIN",
         }
         response = self.client.put(
             f"/api/devices/{self.device.pk}/block/",
             data,
-            format="json",  # ✅ this is crucial!
+            format="json",
         )
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.device.refresh_from_db()

devices/views.py

@@ -63,7 +63,6 @@ class DeviceListCreateAPIView(
             return Response({"message": "MAC address vendor not found."}, status=400)
 
         mac = re.sub(NORMALIZE_MAC_REGEX, "-", mac).upper()
-        request.data["mac"] = mac
 
         return super().create(request, *args, **kwargs)