refactor: centralize credential and param extraction for eks workload list tools

cokerrd · cokerrd · commit 741b2021f768 · 2026-04-29T21:16:31.000+01:00
diff --git a/app/tools/EKSListClustersTool/__init__.py b/app/tools/EKSListClustersTool/__init__.py
@@ -9,6 +9,7 @@
 
 from app.services.eks.eks_client import EKSClient
 from app.tools.tool_decorator import tool
+from app.tools.utils.eks_workload_helper import extract_workload_params
 
 logger = logging.getLogger(__name__)
 
@@ -26,14 +27,6 @@ def _eks_creds(eks: dict) -> dict:
     }
 
 
-def _list_clusters_extract_params(sources: dict[str, dict]) -> dict[str, Any]:
-    eks = sources["eks"]
-    return {
-        "cluster_names": eks.get("cluster_names", []),
-        **_eks_creds(eks),
-    }
-
-
 @tool(
     name="list_eks_clusters",
     source="eks",
@@ -55,7 +48,7 @@ def _list_clusters_extract_params(sources: dict[str, dict]) -> dict[str, Any]:
         "required": ["role_arn"],
     },
     is_available=_eks_available,
-    extract_params=_list_clusters_extract_params,
+    extract_params=extract_workload_params,
 )
 def list_eks_clusters(
     role_arn: str,
diff --git a/app/tools/EKSListDeploymentsTool/__init__.py b/app/tools/EKSListDeploymentsTool/__init__.py
@@ -6,23 +6,13 @@
 from typing import Any, cast
 
 from app.services.eks.eks_k8s_client import build_k8s_clients
-from app.tools.EKSListClustersTool import _eks_creds
 from app.tools.tool_decorator import tool
 from app.tools.utils.availability import eks_available_or_backend
+from app.tools.utils.eks_workload_helper import extract_workload_params
 
 logger = logging.getLogger(__name__)
 
 
-def _list_deployments_extract_params(sources: dict[str, dict]) -> dict[str, Any]:
-    eks = sources["eks"]
-    return {
-        "cluster_name": eks.get("cluster_name", ""),
-        "namespace": eks.get("namespace") or "all",
-        "eks_backend": eks.get("_backend"),
-        **_eks_creds(eks),
-    }
-
-
 @tool(
     name="list_eks_deployments",
     source="eks",
@@ -45,7 +35,7 @@ def _list_deployments_extract_params(sources: dict[str, dict]) -> dict[str, Any]
         "required": ["cluster_name", "namespace", "role_arn"],
     },
     is_available=eks_available_or_backend,
-    extract_params=_list_deployments_extract_params,
+    extract_params=extract_workload_params,
 )
 def list_eks_deployments(
     cluster_name: str,
diff --git a/app/tools/EKSListPodsTool/__init__.py b/app/tools/EKSListPodsTool/__init__.py
@@ -6,23 +6,13 @@
 from typing import Any, cast
 
 from app.services.eks.eks_k8s_client import build_k8s_clients
-from app.tools.EKSListClustersTool import _eks_creds
 from app.tools.tool_decorator import tool
 from app.tools.utils.availability import eks_available_or_backend
+from app.tools.utils.eks_workload_helper import extract_workload_params
 
 logger = logging.getLogger(__name__)
 
 
-def _list_pods_extract_params(sources: dict[str, dict]) -> dict[str, Any]:
-    eks = sources["eks"]
-    return {
-        "cluster_name": eks.get("cluster_name", ""),
-        "namespace": eks.get("namespace") or "all",
-        "eks_backend": eks.get("_backend"),
-        **_eks_creds(eks),
-    }
-
-
 @tool(
     name="list_eks_pods",
     source="eks",
@@ -46,7 +36,7 @@ def _list_pods_extract_params(sources: dict[str, dict]) -> dict[str, Any]:
         "required": ["cluster_name", "namespace", "role_arn"],
     },
     is_available=eks_available_or_backend,
-    extract_params=_list_pods_extract_params,
+    extract_params=extract_workload_params,
 )
 def list_eks_pods(
     cluster_name: str,
diff --git a/app/tools/utils/eks_workload_helper.py b/app/tools/utils/eks_workload_helper.py
@@ -0,0 +1,31 @@
+"""Shared helpers for EKS workload investigation tools"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def _eks_creds(eks: dict) -> dict:
+    """Extract AWS credentials from EKS source"""
+
+    return {
+        "role_arn": eks.get("role_arn", ""),
+        "external_id": eks.get("external_id", ""),
+        "region": eks.get("region", "us-east-1"),
+        "credentials": eks.get("credentials"),
+    }
+
+
+def extract_workload_params(sources: dict[str, dict]) -> dict[str, Any]:
+    """Extract common parameters for workload list operations (pods/deployments)"""
+
+    eks = sources.get("eks")
+    if eks is None:
+        raise ValueError("Sources dictionary must contain an 'eks' key with cluster configuration")
+
+    return {
+        "cluster_name": eks.get("cluster_name", ""),
+        "namespace": eks.get("namespace") or "all",
+        "eks_backend": eks.get("_backend"),
+        **_eks_creds(eks),
+    }
diff --git a/tests/tools/utils/test_eks_workload_helper.py b/tests/tools/utils/test_eks_workload_helper.py
@@ -0,0 +1,60 @@
+"""Tests for eks workload helpers"""
+
+from __future__ import annotations
+
+from app.tools.utils.eks_workload_helper import extract_workload_params
+
+
+def test_extract_basic_params():
+    """Test basic parameter extraction with minimal config"""
+    sources = {"eks": {"cluster_name": "test-cluster", "namespace": "default"}}
+
+    result = extract_workload_params(sources)
+
+    assert result["cluster_name"] == "test-cluster"
+    assert result["namespace"] == "default"
+    assert result["region"] == "us-east-1"
+    assert result["role_arn"] == ""
+    assert result["external_id"] == ""
+    assert result["credentials"] is None
+
+
+def test_namespace_defaults_to_all():
+    """Test namespace defaults to 'all' when not provided"""
+    sources = {"eks": {"cluster_name": "test-cluster"}}
+
+    result = extract_workload_params(sources)
+
+    assert result["namespace"] == "all"
+
+
+def test_handles_all_optional_fields():
+    """Test extraction includes all optional AWS fields"""
+    sources = {
+        "eks": {
+            "cluster_name": "prod-cluster",
+            "role_arn": "arn:aws:iam::123:role/test",
+            "external_id": "external-123",
+            "region": "us-west-2",
+            "credentials": {"access_key": "key123"},
+        }
+    }
+
+    result = extract_workload_params(sources)
+
+    assert result["cluster_name"] == "prod-cluster"
+    assert result["role_arn"] == "arn:aws:iam::123:role/test"
+    assert result["external_id"] == "external-123"
+    assert result["region"] == "us-west-2"
+    assert result["credentials"] == {"access_key": "key123"}
+
+
+def test_missing_eks_raises_error():
+    """Test ValueError when 'eks' key is missing"""
+    sources = {"other": {}}
+
+    try:
+        extract_workload_params(sources)
+        raise AssertionError("Should have raised ValueError")
+    except ValueError as e:
+        assert "must contain an 'eks' key" in str(e)
