tests for prom_metrics that use prometheus rather than local server

Author: memalhot

batchtools/br.py

@@ -21,7 +21,6 @@
 from .prom_metrics import (
     PROMETHEUS_INSTANCE,
     IN_PROGRESS,
-    # record helpers (emit both histogram + counter)
     record_batch_observation,
     record_queue_observation,
     record_wall_observation,
@@ -191,7 +190,6 @@ def run(args: argparse.Namespace):
                     job_name=job_name, wait=True, timeout=args.timeout
                 )
 
-            # Emit metrics if we captured any timing
             if (
                 run_elapsed is not None
                 or queue_wait is not None
@@ -222,8 +220,8 @@ def run(args: argparse.Namespace):
                 IN_PROGRESS.labels(**labels, result=result_phase).dec()
 
                 group = {
-                    "instance": PROMETHEUS_INSTANCE,  # e.g. "ope-test"
-                    "job_name": job_name            # e.g. "job-none-b799d46a6f995a9d8d58b6aff76abefc"
+                    "instance": PROMETHEUS_INSTANCE,  # oc project
+                    "job_name": job_name,  # unique name
                 }
 
                 push_registry_text(grouping_key=group)

batchtools/prom_metrics.py

@@ -1,4 +1,3 @@
-import subprocess
 from datetime import datetime, timezone
 from .helpers import is_on_project
 
@@ -18,9 +17,8 @@
 
 LONG_JOB_BUCKETS = (1, 2, 5, 10, 20, 30, 60, 120, 180, 300, 600, 900, float("inf"))
 
-PUSHGATEWAY_ADDR = os.getenv(
-    "PUSHGATEWAY_ADDR", "pushgateway.ope-test.svc:9091"
-)
+PUSHGATEWAY_ADDR = os.getenv("PUSHGATEWAY_ADDR", "pushgateway.ope-test.svc:9091")
+
 
 def detect_instance() -> str:
     if shutil.which("oc") is None:
@@ -94,6 +92,7 @@ def detect_instance() -> str:
     registry=registry,
 )
 
+
 def now_rfc3339() -> str:
     return datetime.now(timezone.utc).isoformat()
 
@@ -141,6 +140,7 @@ def generate_metrics_text() -> tuple[str, str]:
     payload = generate_latest(registry)
     return payload.decode("utf-8"), CONTENT_TYPE_LATEST
 
+
 def push_registry_text(grouping_key: dict[str, str] | None = None) -> None:
     if not PUSHGATEWAY_ADDR:
         body, _ = generate_metrics_text()

pushgateway/prom.yaml

@@ -48,4 +48,4 @@ spec:
     - name: http
       port: 9091
       targetPort: 9091
-  type: ClusterIP
+  type: ClusterIP

tests/test_prom_metrics.py

@@ -1,19 +1,17 @@
 import batchtools.prom_metrics as pm
-import types
 from datetime import datetime
 from unittest import mock
 
 
 def test_now_rfc3339_parses_with_timezone():
     s = pm.now_rfc3339()
-    # ISO 8601 parseable and includes timezone info (+00:00)
     dt = datetime.fromisoformat(s)
     assert dt.tzinfo is not None
     assert s.endswith("+00:00")
 
 
 def _labels(job="job-x", gpu="none", queue="dummy-localqueue", instance="ns"):
-    return {"job": job, "gpu": gpu, "queue": queue, "instance": instance}
+    return {"job_name": job, "gpu": gpu, "queue": queue, "instance": instance}
 
 
 def _registry_text() -> str:
@@ -78,61 +76,58 @@ def test_generate_metrics_text_returns_valid_payload_and_ctype():
 
 
 def test_push_registry_text_no_url_prints_payload(capsys):
-    # Temporarily clear the push URL so it prints instead of POSTing
-    with mock.patch.object(pm, "PROMETHEUS_PUSH_URL", "", create=True):
+    # check "no PUSHGATEWAY_ADDR" branch
+    with mock.patch.object(pm, "PUSHGATEWAY_ADDR", ""):
         pm.push_registry_text()
+
     out = capsys.readouterr().out
-    assert "PROMETHEUS_PUSH_URL not set" in out
-    assert "# HELP" in out  # the metrics text is printed
+    assert "PROM: PUSHGATEWAY_ADDR not set" in out
+    assert "# HELP" in out
+    assert "# TYPE" in out
 
 
 def test_push_registry_text_posts_success(capsys):
     with (
-        mock.patch.object(
-            pm, "PROMETHEUS_PUSH_URL", "http://example/metrics", create=True
-        ),
-        mock.patch.object(pm.subprocess, "run") as mock_run,
+        mock.patch.object(pm, "PUSHGATEWAY_ADDR", "pushgateway.example:9091"),
+        mock.patch.object(pm, "pushadd_to_gateway") as mock_push,
     ):
-        mock_run.return_value = types.SimpleNamespace(returncode=0)
-
-        pm.push_registry_text()
-        out = capsys.readouterr().out
-        assert "metrics successfully pushed" in out
+        pm.push_registry_text(grouping_key={"instance": "test-ns"})
 
-        # Verify we invoked curl with POST and sent our payload
-        assert mock_run.called
-        args, kwargs = mock_run.call_args
-        argv = args[0]
-        assert "curl" in argv[0]
-        assert "-X" in argv and "POST" in argv
-        assert "http://example/metrics" in argv
+    out = capsys.readouterr().out
+    assert "PROM: metrics pushed to pushgateway=pushgateway.example:9091" in out
 
-        payload = kwargs.get("input", b"").decode("utf-8")
-        assert "# HELP" in payload
-        assert "# TYPE" in payload
+    # verify pushadd_to_gateway was called with expected args
+    mock_push.assert_called_once()
+    args, kwargs = mock_push.call_args
+    assert args[0] == "pushgateway.example:9091"
+    # job and registry are passed as keyword args
+    assert kwargs.get("job") == "batchtools"
+    assert kwargs.get("registry") is pm.registry
+    assert kwargs.get("grouping_key") == {"instance": "test-ns"}
 
 
 def test_push_registry_text_posts_failure(capsys):
     with (
-        mock.patch.object(
-            pm, "PROMETHEUS_PUSH_URL", "http://example/metrics", create=True
-        ),
-        mock.patch.object(pm.subprocess, "run") as mock_run,
+        mock.patch.object(pm, "PUSHGATEWAY_ADDR", "pushgateway.example:9091"),
+        mock.patch.object(pm, "pushadd_to_gateway", side_effect=RuntimeError("boom")),
     ):
-        mock_run.return_value = types.SimpleNamespace(returncode=7)
+        pm.push_registry_text(grouping_key={"instance": "test-ns"})
+
+    out = capsys.readouterr().out
+    assert "PROM: failed to push metrics to pushgateway pushgateway.example:9091" in out
+    assert "boom" in out
 
-        pm.push_registry_text()
-        out = capsys.readouterr().out
-        assert "curl returned nonzero exit 7" in out
 
+def test_push_registry_text_handles_generic_exception(capsys):
+    class WeirdError(Exception):
+        pass
 
-def test_push_registry_text_handles_exception(capsys):
     with (
-        mock.patch.object(
-            pm, "PROMETHEUS_PUSH_URL", "http://example/metrics", create=True
-        ),
-        mock.patch.object(pm.subprocess, "run", side_effect=RuntimeError("boom")),
+        mock.patch.object(pm, "PUSHGATEWAY_ADDR", "pushgateway.example:9091"),
+        mock.patch.object(pm, "pushadd_to_gateway", side_effect=WeirdError("weird")),
     ):
-        pm.push_registry_text()
+        pm.push_registry_text(grouping_key={"instance": "test-ns"})
+
     out = capsys.readouterr().out
-    assert "failed to push metrics via curl" in out
+    assert "PROM: failed to push metrics to pushgateway pushgateway.example:9091" in out
+    assert "weird" in out