Merge pull request #6 from grafana/more-prometheus-tests

Author: sd2k

src/mcp_grafana/client.py

@@ -55,13 +55,13 @@ def __init__(self, url: str, api_key: str | None = None) -> None:
     async def get(self, path: str, params: dict[str, str] | None = None) -> bytes:
         r = await self.c.get(path, params=params)
         if not r.is_success:
-            raise GrafanaError(r.read())
+            raise GrafanaError(r.read().decode())
         return r.read()
 
     async def post(self, path: str, json: dict[str, Any]) -> bytes:
         r = await self.c.post(path, json=json)
         if not r.is_success:
-            raise GrafanaError(r.read())
+            raise GrafanaError(r.read().decode())
         return r.read()
 
     async def list_datasources(self) -> bytes:

src/mcp_grafana/tools/prometheus.py

@@ -22,8 +22,8 @@ async def query_prometheus(
     datasource_uid: str,
     expr: str,
     start_rfc3339: str,
-    end_rfc3339: str,
-    step_seconds: int,
+    end_rfc3339: str | None = None,
+    step_seconds: int | None = None,
     query_type: PrometheusQueryType = "range",
 ) -> DSQueryResponse:
     """
@@ -38,8 +38,13 @@ async def query_prometheus(
     step_seconds: The time series step size in seconds. Ignored if `query_type` is 'instant'.
     query_type: The type of query to use. Either 'range' or 'instant'.
     """
+    if query_type == "range" and (end_rfc3339 is None or step_seconds is None):
+        raise ValueError(
+            "end_rfc3339 and step_seconds must be provided when query_type is 'range'"
+        )
     start = datetime.fromisoformat(start_rfc3339)
-    end = datetime.fromisoformat(end_rfc3339)
+    end = datetime.fromisoformat(end_rfc3339) if end_rfc3339 is not None else start
+    interval_ms = step_seconds * 1000 if step_seconds is not None else None
     query = Query(
         refId="A",
         datasource=DatasourceRef(
@@ -48,7 +53,7 @@ async def query_prometheus(
         ),
         queryType=query_type,
         expr=expr,  # type: ignore
-        intervalMs=step_seconds * 1000,
+        intervalMs=interval_ms,
     )
     response = await grafana_client.query(start, end, [query])
     return DSQueryResponse.model_validate_json(response)

tests/tools/prometheus_test.py

@@ -2,7 +2,8 @@
 
 import pytest
 
-from mcp_grafana.grafana_types import LabelMatcher, Selector
+from mcp_grafana.client import GrafanaError
+from mcp_grafana.grafana_types import DSQueryResponse, LabelMatcher, Selector
 from mcp_grafana.tools.prometheus import (
     query_prometheus,
     get_prometheus_metric_metadata,
@@ -16,29 +17,104 @@
 pytestmark = mark_integration
 
 
-@pytest.mark.parametrize("step_seconds", [15, 60, 300])
-async def test_query_prometheus_range(step_seconds):
-    end = datetime.now()
-    start = end - timedelta(minutes=10)
+class TestPrometheusQueries:
+    @pytest.mark.parametrize("step_seconds", [15, 60, 300])
+    async def test_query_prometheus_range(self, step_seconds):
+        end = datetime.now()
+        start = end - timedelta(minutes=10)
 
-    results = await query_prometheus(
-        "robustperception",
-        start_rfc3339=start.isoformat(),
-        end_rfc3339=end.isoformat(),
-        step_seconds=step_seconds,
-        expr="node_load1",
-        query_type="range",
-    )
-    query_result = results.results["A"]
-    assert query_result.status == 200
-    assert (
-        query_result.frames[0].data.values[0][0] - start.timestamp() * 1000
-        < step_seconds * 1000
-    )
-    assert (
-        len(query_result.frames[0].data.values[0])
-        == (end - start).total_seconds() // step_seconds + 1
-    )
+        results = await query_prometheus(
+            "robustperception",
+            start_rfc3339=start.isoformat(),
+            end_rfc3339=end.isoformat(),
+            step_seconds=step_seconds,
+            expr="node_load1",
+            query_type="range",
+        )
+        query_result = results.results["A"]
+        assert query_result.status == 200
+        assert (
+            query_result.frames[0].data.values[0][0] - start.timestamp() * 1000
+            < step_seconds * 1000
+        )
+        assert (
+            len(query_result.frames[0].data.values[0])
+            == (end - start).total_seconds() // step_seconds + 1
+        )
+
+    async def test_query_prometheus_instant(self):
+        timestamp = datetime.now()
+        results = await query_prometheus(
+            "robustperception",
+            start_rfc3339=timestamp.isoformat(),
+            expr="up",
+            query_type="instant",
+        )
+        query_result = results.results["A"]
+
+        # Verify the response
+        assert query_result.status == 200
+        assert len(query_result.frames) > 0
+        # Instant queries should return a single data point
+        assert len(query_result.frames[0].data.values[0]) == 1
+        # Verify we have both time and value columns
+        assert len(query_result.frames[0].data.values) == 2
+        # Verify the timestamp is close to our request time (within 60 seconds,
+        # we don't know what the scrape interval is).
+        assert (
+            abs(query_result.frames[0].data.values[0][0] - timestamp.timestamp() * 1000)
+            < 60000
+        )
+
+    async def test_query_prometheus_range_missing_params(self):
+        start = datetime.now()
+        with pytest.raises(
+            ValueError, match="end_rfc3339 and step_seconds must be provided"
+        ):
+            await query_prometheus(
+                "robustperception",
+                start_rfc3339=start.isoformat(),
+                expr="node_load1",
+                query_type="range",
+            )
+
+    async def test_query_prometheus_invalid_query(self):
+        end = datetime.now()
+        start = end - timedelta(minutes=10)
+
+        try:
+            results = await query_prometheus(
+                "robustperception",
+                start_rfc3339=start.isoformat(),
+                end_rfc3339=end.isoformat(),
+                step_seconds=60,
+                expr="invalid_metric{",  # Invalid PromQL syntax
+                query_type="range",
+            )
+        except Exception as e:
+            assert isinstance(e, GrafanaError)
+            results = DSQueryResponse.model_validate_json(str(e))
+        query_result = results.results["A"]
+        assert query_result.status == 400
+        assert query_result.error is not None
+        assert "parse error" in query_result.error
+
+    async def test_query_prometheus_empty_result(self):
+        end = datetime.now()
+        start = end - timedelta(minutes=10)
+
+        results = await query_prometheus(
+            "robustperception",
+            start_rfc3339=start.isoformat(),
+            end_rfc3339=end.isoformat(),
+            step_seconds=60,
+            expr='metric_that_does_not_exist{label="value"}',
+            query_type="range",
+        )
+        query_result = results.results["A"]
+        assert query_result.status == 200
+        assert len(query_result.frames) == 1
+        assert len(query_result.frames[0].data.values) == 0
 
 
 async def test_get_prometheus_metric_metadata():