Add update and delete methods for Kentik device API

gso/services/kentik_client.py

@@ -4,13 +4,29 @@ import logging
 from typing import Any
 
 import requests
+from pydantic import BaseModel
+from requests import Response
 
-from gso.products.product_types.router import Router
+from gso.products.product_blocks.site import SiteTier
 from gso.settings import load_oss_params
 
 logger = logging.getLogger(__name__)
 
 
+class NewKentikDevice(BaseModel):
+    """API model for creating a new device in Kentik."""
+
+    device_name: str
+    device_description: str
+    sending_ips: list[str]
+    site_id: int
+    site_tier: SiteTier
+    device_snmp_ip: str
+    device_bgp_flowspec: bool
+    device_bgp_neighbor_ip: str
+    device_bgp_neighbor_ip6: str
+
+
 class KentikClient:
     """The client for Kentik that interacts with an external instance."""
 
@@ -23,19 +39,21 @@ class KentikClient:
             "Content-Type": "application/json",
         }
 
-    def _send_request(self, method, endpoint: str, data: dict[str, Any] | None = None) -> dict[str, Any]:
+    def _send_request(self, method, endpoint: str, data: dict[str, Any] | None = None) -> Response:
         url = self.config.api_base + endpoint
+        logger.debug("Kentik - Sending request", extra={"method": method, "endpoint": url, "form_data": data})
         result = requests.request(method, url, json=data, headers=self.headers)
+        logger.debug("Kentik - Received response", extra=result.__dict__)
 
-        return result.json()
+        return result
 
     def get_devices(self) -> list[dict[str, Any]]:
         """List all devices in Kentik."""
-        return [self._send_request("GET", "v5/devices")]
+        return [self._send_request("GET", "v5/devices").json()]
 
     def get_device(self, device_id: str) -> dict[str, Any]:
         """Get a device by ID."""
-        return self._send_request("GET", f"v5/device/{device_id}")
+        return self._send_request("GET", f"v5/device/{device_id}").json()
 
     def get_device_by_name(self, device_name: str) -> dict[str, Any]:
         """Fetch a device in Kentik by its :term:`FQDN`.
@@ -53,11 +71,11 @@ class KentikClient:
 
     def get_sites(self) -> list[dict[str, Any]]:
         """Get a list of all available sites in Kentik."""
-        return self._send_request("GET", "v5/sites")["sites"]
+        return self._send_request("GET", "v5/sites").json()["sites"]
 
     def get_site(self, site_id: str) -> dict[str, Any]:
         """Get a site by ID."""
-        return self._send_request("GET", f"v5/site/{site_id}")
+        return self._send_request("GET", f"v5/site/{site_id}").json()
 
     def get_site_by_name(self, site_slug: str) -> dict[str, Any]:
         """Get a Kentik site by its name.
@@ -75,7 +93,7 @@ class KentikClient:
 
     def get_plans(self) -> list[dict[str, Any]]:
         """Get all Kentik plans available."""
-        return self._send_request("GET", "v5/plans")["plans"]
+        return self._send_request("GET", "v5/plans").json()["plans"]
 
     def get_plan_by_name(self, plan_name: str) -> dict[str, Any]:
         """Get a Kentik plan by its name.
@@ -91,31 +109,38 @@ class KentikClient:
 
         return {}
 
-    def create_device(self, router: Router) -> dict[str, Any]:
+    def create_device(self, device: NewKentikDevice) -> dict[str, Any]:
         """Add a new device to Kentik."""
-        site_id = self.get_site_by_name(router.router.router_site.site_name)["id"]
-        plan_id = self.get_plan_by_name(self.config.billing_plans[router.router.router_site.site_tier])
+        plan_id = self.get_plan_by_name(self.config.billing_plans[device.site_tier])["id"]
         request_body = {
             "device": {
-                # Route selection missing
-                "deviceName": router.router.router_fqdn,
-                "deviceSubtype": self.config.device_type,
-                "sendingIps": [str(router.router.router_lo_ipv4_address)],
-                "deviceSampleRate": self.config.sample_rate,
-                "planId": plan_id,
-                "siteId": site_id,
-                "deviceSnmpIp": str(router.router.router_lo_ipv4_address),
-                "deviceSnmpCommunity": self.config.snmp_community,
-                "deviceBgpType": self.config.bgp_type,
-                "deviceBgpFlowspec": False,
-                "deviceBgpNeighborIp": str(router.router.router_lo_ipv4_address),
-                "deviceBgpNeighborIp6": str(router.router.router_lo_ipv6_address),
-                "deviceBgpNeighborAsn": str(self.config.ASN),
-                "deviceBgpPassword": self.config.md5_password,
+                **device.model_dump(exclude=set("site_tier")),
+                "device_type": self.config.device_type,
+                "device_subtype": self.config.device_type,
+                "minimize_snmp": self.config.minimize_snmp,
+                "device_sample_rate": self.config.sample_rate,
+                "plan_id": plan_id,
+                "device_snmp_community": self.config.snmp_community,
+                "device_bgp_type": self.config.bgp_type,
+                "bgp_lookup_strategy": self.config.bgp_lookup_strategy,
+                "device_bgp_neighbor_asn": str(self.config.ASN),
+                "device_bgp_password": self.config.md5_password,
             }
         }
 
-        return self._send_request("POST", "v5/device", request_body)
+        return self._send_request("POST", "v5/device", request_body).json()
+
+    def update_device(self, device_id: str, updated_device: dict[str, Any]) -> dict[str, Any]:
+        """Update an existing device in Kentik."""
+        return self._send_request("PUT", f"v5/device/{device_id}", updated_device).json()
+
+    def remove_device(self, device_id: str, *, archive: bool = True) -> None:
+        """Remove a device from Kentik.
+
+        :param str device_id: The Kentik internal ID of the device that is to be removed.
+        :param bool archive: Archive the device instead of completely deleting it.
+        """
+        if not archive:
+            self._send_request("DELETE", f"v5/device/{device_id}")
 
-    def remove_device(self, hostname: str) -> None:
-        """Remove a device from Kentik."""
+        self._send_request("DELETE", f"v5/device/{device_id}")