FabledSteward/tests/integration/test_prometheus_scrape.py

"""Integration: a Prometheus scrape writes each target's outcome to its row,
its kept series to the metric tables (a counter as a rate), and its own
is_up, which an alert rule can watch (#5668). A series a good scrape no
longer reports loses its newest value (#5689); a failed scrape keeps them.
Requires STEWARD_DATABASE_URL."""
from __future__ import annotations

import asyncio
import importlib
import os
import uuid

import pytest

pytestmark = pytest.mark.integration

TEXT = """\
# TYPE fa_backup_target_state gauge
fa_backup_target_state{target="Steward",state="STATE"} 1
# TYPE fa_requests_total counter
fa_requests_total COUNT
# TYPE noise gauge
noise 1
"""


@pytest.fixture
def app():
    if not os.environ.get("STEWARD_DATABASE_URL"):
        pytest.skip("needs Postgres")
    from steward.app import create_app
    return create_app(testing=False)


def _prom(sub):
    """A submodule of the prometheus plugin as the app loaded it: it is on by
    default, so its models are already mapped under the loader's name, and
    importing plugins.prometheus.models would define the table a second time."""
    from steward.core.plugin_manager import get_loaded_plugin_modules
    loaded = get_loaded_plugin_modules().get("prometheus")
    return importlib.import_module(f"{loaded.__name__ if loaded else 'plugins.prometheus'}.{sub}")


def test_a_scrape_records_series_rates_and_a_down_target_fires_its_rule(app):
    from sqlalchemy import select

    from steward.alerts import service as alert_svc
    from steward.core import prometheus
    from steward.core.alerts import record_metrics
    from steward.core.crypto import encrypt_secret
    from steward.models.alerts import AlertState, AlertStateEnum
    from steward.models.metrics import PluginMetric, PluginMetricLatest
    from steward.models.users import User, UserRole

    scheduler = _prom("scheduler")
    PrometheusTarget = _prom("models").PrometheusTarget
    ScrapeError = _prom("scraper").ScrapeError

    tag = uuid.uuid4().hex[:8]
    up, down = f"arch-{tag}", f"gone-{tag}"
    counts = iter([1_000, 1_011])
    states = iter(["late", "fresh"])
    tokens = []

    async def fake_scrape(url, token, verify_tls):
        tokens.append(token)
        if "gone" in url:
            raise ScrapeError("Couldn't connect: connection refused.")
        return prometheus.parse(TEXT.replace("COUNT", str(next(counts))).replace("STATE", next(states)))

    async def latest(name_prefix):
        async with app.db_sessionmaker() as s:
            rows = (await s.execute(select(PluginMetricLatest).where(
                PluginMetricLatest.source_module == "prometheus",
                PluginMetricLatest.resource_name.startswith(name_prefix)))).scalars()
            return {(r.resource_name, r.metric_name): r.value for r in rows}

    async def go():
        async with app.db_sessionmaker() as s:
            async with s.begin():
                user = User(username=f"pr-{tag}", email=f"pr-{tag}@test", password_hash="x",
                            role=UserRole.admin)
                s.add_all([user,
                           PrometheusTarget(name=up, url=f"http://{up}/metrics",
                                            token=encrypt_secret("scrape-token"), keep=["fa_*"]),
                           PrometheusTarget(name=down, url=f"http://{down}/metrics")])
                await s.flush()
                await record_metrics(s, [PluginMetric(source_module="prometheus", resource_name=down,
                                                      metric_name="is_up", value=1.0)])
                await s.flush()
                fields = await alert_svc.rule_fields(
                    s, name=f"down {tag}", source_module="prometheus", resource_name=down,
                    metric_name="is_up", operator="<", threshold=1)
                rule = await alert_svc.add_rule(s, fields, created_by=user.id)
            rule_id = rule.id

        await scheduler.scrape_due(app, scrape=fake_scrape)
        first = await latest(up)
        # Every target was just scraped, so nothing is due on the next tick.
        await scheduler.scrape_due(app, scrape=fake_scrape)
        assert len(tokens) == 2
        async with app.db_sessionmaker() as s:
            async with s.begin():
                for t in (await s.execute(select(PrometheusTarget).where(
                        PrometheusTarget.name.in_([up, down])))).scalars():
                    t.last_scraped_at = t.last_scraped_at.replace(year=2000)
        await asyncio.sleep(1.1)
        await scheduler.scrape_due(app, scrape=fake_scrape)
        second = await latest(up)
        async with app.db_sessionmaker() as s:
            rows = {t.name: t for t in (await s.execute(select(PrometheusTarget).where(
                PrometheusTarget.name.in_([up, down])))).scalars()}
            state = (await s.execute(select(AlertState).where(
                AlertState.rule_id == rule_id))).scalar_one()
        async with app.db_sessionmaker() as s:
            async with s.begin():
                for t in rows.values():
                    await s.delete(await s.get(PrometheusTarget, t.id))
        return first, second, rows, state, await latest(down)

    first, second, rows, state, down_latest = asyncio.run(go())
    assert "scrape-token" in tokens
    assert first[(f"{up}:target=Steward,state=late", "fa_backup_target_state")] == 1.0
    # The second scrape no longer reports "late": its newest value is forgotten.
    assert (f"{up}:target=Steward,state=late", "fa_backup_target_state") not in second
    assert second[(f"{up}:target=Steward,state=fresh", "fa_backup_target_state")] == 1.0
    assert (up, "fa_requests_per_sec") not in first                  # a counter needs two scrapes
    assert first[(up, "is_up")] == 1.0 and first[(up, "series_kept")] == 2.0
    assert not any(m == "noise" for _, m in first)                   # not in the keep list
    # 11 requests over the 1.1 s or more between the scrapes.
    assert 0 < second[(up, "fa_requests_per_sec")] <= 10.0
    assert rows[up].last_error is None and rows[up].last_kept == 2 and rows[up].last_dropped == 1
    assert rows[down].last_error == "Couldn't connect: connection refused."
    assert rows[down].last_ok_at is None
    assert down_latest[(down, "is_up")] == 0.0
    assert state.state == AlertStateEnum.firing