From 55be22cb1097db145924f7e7f88ed36bb4e21cfa Mon Sep 17 00:00:00 2001
From: Lars Michelsen <lm@checkmk.com>
Date: Fri, 14 Jun 2024 23:11:35 +0200
Subject: [PATCH] Cleanup no-untyped-def in agent_prometheus

Change-Id: I539d8db2c66912f08d04118c2fd125af499fa8fb
---
 cmk/plugins/lib/prometheus.py                  |  8 ++++++--
 .../special_agents/agent_prometheus.py         | 18 +++++++++++-------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/cmk/plugins/lib/prometheus.py b/cmk/plugins/lib/prometheus.py
index f39dcbc7843..def2de4ad2b 100755
--- a/cmk/plugins/lib/prometheus.py
+++ b/cmk/plugins/lib/prometheus.py
@@ -10,7 +10,11 @@
 
 from cmk.utils import password_store
 
-from cmk.special_agents.v0_unstable.request_helper import create_api_connect_session, parse_api_url
+from cmk.special_agents.v0_unstable.request_helper import (
+    ApiSession,
+    create_api_connect_session,
+    parse_api_url,
+)
 
 
 class ConnectionConfig(TypedDict):
@@ -49,7 +53,7 @@ def extract_connection_args(config):
     return connection_args
 
 
-def generate_api_session(connection_options):
+def generate_api_session(connection_options: dict) -> ApiSession:
     return create_api_connect_session(
         connection_options["api_url"],
         auth=connection_options.get("auth"),
diff --git a/cmk/plugins/prometheus/special_agents/agent_prometheus.py b/cmk/plugins/prometheus/special_agents/agent_prometheus.py
index 07617e76b96..2977c8b7379 100644
--- a/cmk/plugins/prometheus/special_agents/agent_prometheus.py
+++ b/cmk/plugins/prometheus/special_agents/agent_prometheus.py
@@ -26,6 +26,7 @@
     extract_connection_args,
     generate_api_session,
 )
+from cmk.special_agents.v0_unstable.request_helper import ApiSession
 
 LOGGER = logging.getLogger()  # root logger for now
 
@@ -80,7 +81,7 @@ def parse_pod_name(labels: dict[str, str], prepend_namespace: bool = False) -> s
 
 
 class CAdvisorExporter:
-    def __init__(self, api_client, options) -> None:  # type: ignore[no-untyped-def]
+    def __init__(self, api_client: "PrometheusAPI", options: dict) -> None:
         self.api_client = api_client
         self.container_name_option = options.get("container_id", "short")
         self.pod_containers: dict = {}
@@ -192,7 +193,7 @@ def memory_pod_summary(self, _group_element: str) -> list[dict[str, dict[str, An
     def _retrieve_pods_memory_summary(
         self, memory_info: list[tuple[str, str]]
     ) -> tuple[dict[str, dict[str, Any]], dict[str, str]]:
-        result: dict[str, dict[str, str | dict[str, str]]] = {}
+        result: dict[str, dict[str, str | dict[str, Any]]] = {}
         associations = {}
         for memory_stat, promql_query in memory_info:
             promql_query = promql_query.replace("{namespace_filter}", self._namespace_query_part())
@@ -400,8 +401,8 @@ def __init__(self, raw_response: dict[str, Any]) -> None:
         self.labels = raw_response["metric"]
         self.internal_values = raw_response["value"]
 
-    def label_value(self, key: str, default=None) -> str:  # type: ignore[no-untyped-def]
-        return self.labels.get(key, default)
+    def label_value(self, key: str) -> str:
+        return self.labels.get(key)
 
     def value(self, default_value: float | int | None = None, as_string: bool = False) -> float:
         try:
@@ -563,7 +564,7 @@ class PrometheusAPI:
     Realizes communication with the Prometheus API
     """
 
-    def __init__(self, session) -> None:  # type: ignore[no-untyped-def]
+    def __init__(self, session: ApiSession) -> None:
         self.session = session
 
     def perform_specified_promql_queries(
@@ -737,7 +738,7 @@ class ApiData:
     Server & the Prometheus Exporters
     """
 
-    def __init__(self, api_client, exporter_options) -> None:  # type: ignore[no-untyped-def]
+    def __init__(self, api_client: "PrometheusAPI", exporter_options: dict) -> None:
         self.api_client = api_client
         self.prometheus_server = PrometheusServer(api_client)
         if "cadvisor" in exporter_options:
@@ -746,7 +747,10 @@ def __init__(self, api_client, exporter_options) -> None:  # type: ignore[no-unt
         if "node_exporter" in exporter_options:
 
             def get_promql(promql_expression: str) -> list[PromQLMetric]:
-                return api_client.perform_multi_result_promql(promql_expression).promql_metrics
+                result = api_client.perform_multi_result_promql(promql_expression)
+                if result is None:
+                    raise ApiError("Missing PromQL result for %s" % promql_expression)
+                return result.promql_metrics
 
             self.node_exporter = NodeExporter(get_promql)