broad-announce/scripts/m11_lib.py

"""
m11_lib.py — shared assertion library for M11 smoke scripts.
Used by m11_smoke.py.

All Prometheus queries use URL-encoding-safe Python urllib.
"""
import urllib.request
import urllib.parse
import json

PROM = "http://localhost:9090"
PROM_TIMEOUT = 10  # seconds


def scrape(query: str) -> list[dict]:
    """
    Run a Prometheus query and return the result vector.
    Returns [] on error or no data.
    """
    url = f"{PROM}/api/v1/query?query={urllib.parse.quote(query, safe='')}"
    try:
        with urllib.request.urlopen(url, timeout=PROM_TIMEOUT) as r:
            data = r.read()
        d = json.loads(data)
        if d.get("status") != "success":
            return []
        return d.get("data", {}).get("result", [])
    except Exception:
        return []


def assert_grpc_rate_near(target: float, tolerance: float = 0.10,
                          window_seconds: int = 30) -> float:
    """
    Assert the gRPC ingest rate is within tolerance of target.
    target: expected alerts/sec cluster-wide
    tolerance: fraction (0.10 = ±10%)
    Returns the actual rate, or raises AssertionError.

    Query: sum(rate(ba_ingestd_alerts_received_total{transport="grpc",result="ok"}[window]))
    """
    query = (
        f'sum(rate(ba_ingestd_alerts_received_total{{transport="grpc",result="ok"}}[{window_seconds}s]))'
    )
    results = scrape(query)
    if not results:
        raise AssertionError(
            f"gRPC rate query returned no data. Is ingestd gRPC server up and receiving load? "
            f"(query: {query})"
        )
    rate = float(results[0]["value"][1])
    min_rate = target * (1 - tolerance)
    max_rate = target * (1 + tolerance)
    if rate < min_rate:
        raise AssertionError(
            f"gRPC rate {rate:.0f}/s is below target {target:.0f}/s "
            f"(tolerance ±{tolerance*100:.0f}%, min allowed: {min_rate:.0f}/s, "
            f"window={window_seconds}s)"
        )
    return rate


def assert_grpc_p99_under(threshold_ms: float, window_seconds: int = 60) -> float:
    """
    Assert gRPC ack p99 latency is under threshold (in milliseconds).
    Uses histogram_quantile over the given window.
    Returns the actual p99 in milliseconds, or raises AssertionError.

    Query: histogram_quantile(0.99, rate(ba_ingestd_grpc_ack_latency_seconds_bucket[window]))
    """
    query = (
        f'histogram_quantile(0.99, '
        f'rate(ba_ingestd_grpc_ack_latency_seconds_bucket[{window_seconds}s]))'
    )
    results = scrape(query)
    if not results:
        raise AssertionError(
            f"gRPC p99 query returned no data. Is ingestd gRPC server up? "
            f"(query: {query})"
        )
    p99_seconds = float(results[0]["value"][1])
    p99_ms = p99_seconds * 1000
    if p99_ms > threshold_ms:
        raise AssertionError(
            f"gRPC ack p99 {p99_ms:.1f}ms exceeds threshold {threshold_ms:.1f}ms "
            f"(window={window_seconds}s)"
        )
    return p99_ms


def assert_dlq_count_equals(expected: int, window_seconds: int) -> int:
    """
    Assert the DLQ row count (delta over window) equals expected.
    Returns the actual count.
    """
    query = (
        f'increase(ba_deliverd_dlq_total[{window_seconds}s])'
    )
    results = scrape(query)
    if not results:
        return 0  # No data → assume 0
    count = float(results[0]["value"][1])
    if abs(count - expected) > 0.5:
        raise AssertionError(
            f"DLQ count {count:.0f} over last {window_seconds}s does not equal expected {expected}"
        )
    return int(count)


def get_prometheus_targets() -> dict[str, str]:
    """
    Return a dict mapping service name → health status from Prometheus target health API.
    "1" = healthy, "0" = down, other = unknown state.
    """
    url = f"{PROM}/api/v1/targets?state=active"
    try:
        with urllib.request.urlopen(url, timeout=PROM_TIMEOUT) as r:
            data = json.loads(r.read())
        out = {}
        for t in data.get("data", {}).get("activeTargets", []):
            labels = t.get("labels", {})
            job = labels.get("job", "?")
            # Use the job name (e.g. "ingestd", "loadgen-grpc-1") as the key
            out[job] = "1" if t.get("health") == "up" else str(t.get("health", "0"))
        return out
    except Exception:
        return {}


def get_grpc_streams_active() -> int:
    """
    Return the current number of active gRPC streams (ba_ingestd_grpc_streams_active gauge).
    Returns 0 if no data.
    """
    query = 'ba_ingestd_grpc_streams_active'
    results = scrape(query)
    if not results:
        return 0
    return int(float(results[0]["value"][1]))


def get_grpc_rate_limited_total() -> int:
    """
    Return the cumulative ba_ingestd_grpc_rate_limited_total counter.
    Returns 0 if no data.
    """
    query = 'ba_ingestd_grpc_rate_limited_total'
    results = scrape(query)
    if not results:
        return 0
    return int(float(results[0]["value"][1]))