Apply archiving license to Kentik device when terminating a PE router

gso/oss-params-example.json

@@ -113,6 +113,7 @@
     "device_type": "router",
     "minimize_snmp": false,
     "placeholder_license_key": "placeholder license",
+    "archive_license_key": "archiving license",
     "sample_rate": 100,
     "bgp_type": "device",
     "bgp_lookup_strategy": "lu_global_fallback",

gso/settings.py

@@ -195,6 +195,7 @@ class KentikParams(BaseSettings):
     device_type: str
     minimize_snmp: bool
     placeholder_license_key: str
+    archive_license_key: str
     sample_rate: int
     bgp_type: str
     bgp_lookup_strategy: str

gso/workflows/router/terminate_router.py

@@ -3,11 +3,13 @@
 import ipaddress
 import json
 import logging
+from typing import Any
 
 from orchestrator.forms import FormPage
 from orchestrator.forms.validators import Label
 from orchestrator.targets import Target
 from orchestrator.types import FormGenerator, SubscriptionLifecycle, UUIDstr
+from orchestrator.utils.errors import ProcessFailureError
 from orchestrator.utils.json import json_dumps
 from orchestrator.workflow import StepList, begin, conditional, done, step, workflow
 from orchestrator.workflows.steps import (
@@ -21,9 +23,11 @@ from orchestrator.workflows.utils import wrap_modify_initial_input_form
 from gso.products.product_blocks.router import RouterRole
 from gso.products.product_types.router import Router
 from gso.services import infoblox, lso_client
+from gso.services.kentik_client import KentikClient
 from gso.services.librenms_client import LibreNMSClient
 from gso.services.lso_client import execute_playbook, lso_interaction
 from gso.services.netbox_client import NetboxClient
+from gso.settings import load_oss_params
 from gso.utils.helpers import generate_inventory_for_active_routers
 from gso.utils.shared_enums import Vendor
 from gso.utils.types import TTNumber
@@ -232,10 +236,28 @@ def remove_pe_from_all_p_real(subscription: Router, callback_route: str, tt_numb
 
 
 @step("Remove Device from Librenms")
-def remove_device_from_librenms(subscription: Router) -> dict[str, Router]:
+def remove_device_from_librenms(subscription: Router) -> None:
     """Remove the device from LibreNMS."""
     LibreNMSClient().remove_device(subscription.router.router_fqdn)
-    return {"subscription": subscription}
+
+
+@step("Apply the archiving license in Kentik")
+def kentik_apply_archive_license(subscription: Router) -> dict[str, dict[str, Any]]:
+    """Apply the archiving license to a PE router in Kentik.
+
+    This includes setting the flow rate to one flow per second.
+    """
+    kentik_client = KentikClient()
+    kentik_archive_plan_id = kentik_client.get_plan_by_name(load_oss_params().KENTIK.archive_license_key)["id"]
+    kentik_device = kentik_client.get_device_by_name(subscription.router.router_fqdn)
+    if "id" not in kentik_device:
+        msg = "Failed to find Kentik device by name"
+        raise ProcessFailureError(msg, details=kentik_device)
+
+    updated_device = {"device": {"plan_id": kentik_archive_plan_id, "device_sample_rate": 1}}
+    kentik_device = kentik_client.update_device(kentik_device["id"], updated_device)
+
+    return {"kentik_device": kentik_device}
 
 
 @workflow(
@@ -272,6 +294,7 @@ def terminate_router() -> StepList:
         >> run_config_steps(lso_interaction(remove_config_from_router_real))
         >> router_is_nokia(remove_device_from_netbox)
         >> remove_device_from_librenms
+        >> router_is_pe(kentik_apply_archive_license)
         >> set_status(SubscriptionLifecycle.TERMINATED)
         >> resync
         >> done

test/workflows/router/test_terminate_router.py

@@ -4,6 +4,7 @@ import pytest
 
 from gso.products import Router
 from gso.products.product_blocks.router import RouterRole
+from test.services.conftest import MockedKentikClient
 from test.workflows import assert_complete, assert_lso_interaction_success, extract_state, run_workflow
 
 
@@ -13,9 +14,11 @@ from test.workflows import assert_complete, assert_lso_interaction_success, extr
 @patch("gso.services.lso_client._send_request")
 @patch("gso.workflows.router.terminate_router.NetboxClient.delete_device")
 @patch("gso.workflows.router.terminate_router.infoblox.delete_host_by_ip")
+@patch("gso.workflows.router.terminate_router.KentikClient")
 @patch("gso.workflows.router.terminate_router.LibreNMSClient.remove_device")
 def test_terminate_pe_router_full_success(
     mock_librenms_remove_device,
+    mock_kentik_client,
     mock_delete_host_by_ip,
     mock_delete_device,
     mock_execute_playbook,
@@ -37,6 +40,7 @@ def test_terminate_pe_router_full_success(
         lso_interaction_count += 2
     if update_ibgp_mesh:
         lso_interaction_count += 4
+    mock_kentik_client.return_value = MockedKentikClient
     #  Run workflow
     initial_router_data = [{"subscription_id": product_id}, router_termination_input_form_data]
     result, process_stat, step_log = run_workflow("terminate_router", initial_router_data)