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

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_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
+
+omada_client = Omada()
 
 
 @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:
-        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 []
+    return omada_client.get_existing_omada_devices()
 
 
 @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:
-        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}")
+    omada_client.add_new_devices_to_omada(new_devices)
 
 
 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