vulncheck/tests/test_dashboard_stats_queries.py

"""Dashboard stats are one aggregate pass — run: python tests/test_dashboard_stats_queries.py

/vulnerabilities/reports/dashboard issued eighteen separate COUNTs over the
same vulnerabilities-joined-to-active-assets scan (four severities, exploitable,
avg, distinct assets, oldest row, seven one-day history counts, two totals, two
asset counts). Every one of them re-scanned the table, and because the endpoint
is `async def` with a blocking Session they run strictly back to back while the
rest of the dashboard's requests wait behind them.

Pinned here: a flat, small statement count and the exact numbers the eighteen
counts produced. Plus the list endpoint's `with_total=false`, which drops the
second full window-function pass the dashboard widgets never render.
"""
import os
import sys
from datetime import datetime, timedelta

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.environ.setdefault("JWT_SECRET_KEY", "test-key-test-key-test-key-test-key")

from sqlalchemy import create_engine, event  # noqa: E402
from sqlalchemy.orm import sessionmaker  # noqa: E402

import app.models  # noqa: F401,E402  (registers every mapper)
from app.models.base import Base  # noqa: E402
from app.models.asset import Asset, AssetSource, AssetStatus  # noqa: E402
from app.models.vulnerability import (  # noqa: E402
    Vulnerability, VulnerabilitySeverity, VulnerabilityStatus,
)
from app.routers.vulnerabilities import (  # noqa: E402
    get_dashboard_statistics, list_vulnerabilities,
)

SEVERITIES = [VulnerabilitySeverity.critical, VulnerabilitySeverity.high,
              VulnerabilitySeverity.medium, VulnerabilitySeverity.low]


def _db(hosts=10, per_host=8):
    engine = create_engine("sqlite://")
    Base.metadata.create_all(engine)
    db = sessionmaker(bind=engine)()
    now = datetime.now()
    for i in range(hosts):
        asset = Asset(hostname=f"host-{i}", ip_address="10.0.0.1",
                      source=AssetSource.MANUAL, operating_system="Ubuntu",
                      # One host is decommissioned — its findings must not count.
                      status=AssetStatus.DECOMMISSIONED if i == 0 else AssetStatus.ACTIVE,
                      last_scan=now if i % 2 == 0 else None)
        db.add(asset)
        db.flush()
        for j in range(per_host):
            n = i * per_host + j
            db.add(Vulnerability(
                cve_id=f"CVE-2024-{n:05d}", asset_id=asset.id,
                severity=SEVERITIES[j % 4],
                status=VulnerabilityStatus.patched if j == 0 else VulnerabilityStatus.open,
                cvss_score=float(j),
                exploit_available=(j % 3 == 0),
                published_date=now - timedelta(days=j),
                detected_at=now - timedelta(days=j % 5),
                sources='["wazuh"]',
            ))
    db.commit()
    return engine, db


class _User:
    pass


def _stats(db):
    return get_dashboard_statistics(db=db, current_user=_User())


def test_statement_count_is_flat_and_small():
    counts = {}
    for hosts in (10, 100):
        engine, db = _db(hosts=hosts)
        stmts = []
        event.listen(engine, "before_cursor_execute",
                     lambda conn, cur, stmt, *a: stmts.append(stmt))
        _stats(db)
        counts[hosts] = len(stmts)
    assert counts[10] == counts[100], f"statements grew with the fleet: {counts}"
    assert counts[100] <= 4, f"expected the aggregate pass, got {counts[100]} statements"
    print(f"✅ {counts[100]} statements for 100 hosts, same as for 10 (was 18)")


def test_numbers_match_the_per_count_semantics():
    """Same values the eighteen separate COUNTs produced."""
    engine, db = _db(hosts=10, per_host=8)
    s = _stats(db)

    active = db.query(Asset).filter(Asset.status == AssetStatus.ACTIVE).all()
    active_ids = {a.id for a in active}
    rows = [v for v in db.query(Vulnerability).all() if v.asset_id in active_ids]
    open_rows = [v for v in rows if v.status == VulnerabilityStatus.open]

    assert s["total_vulnerabilities"] == len(rows)
    assert s["open_vulnerabilities"] == len(open_rows)
    for key, sev in (("critical_count", VulnerabilitySeverity.critical),
                     ("high_count", VulnerabilitySeverity.high),
                     ("medium_count", VulnerabilitySeverity.medium),
                     ("low_count", VulnerabilitySeverity.low)):
        assert s[key] == len([v for v in open_rows if v.severity == sev]), key
    assert s["exploitable_count"] == len([v for v in open_rows if v.exploit_available])
    assert s["avg_cvss_score"] == round(
        sum(v.cvss_score for v in open_rows) / len(open_rows), 2)
    assert s["affected_assets"] == len({v.asset_id for v in open_rows})
    assert s["oldest_vulnerability_days"] == \
        (datetime.now() - min(v.detected_at for v in open_rows)).days
    assert s["total_assets"] == db.query(Asset).count()
    assert s["scanned_assets"] == db.query(Asset).filter(Asset.last_scan.isnot(None)).count()
    print("✅ every aggregate matches the per-count result, decommissioned host excluded")


def test_seven_day_history_still_has_seven_days():
    engine, db = _db()
    s = _stats(db)
    hist = s["severity_history"]
    assert len(hist) == 7
    now = datetime.now()
    assert [h["day"] for h in hist] == \
        [(now - timedelta(days=i)).strftime("%a") for i in range(6, -1, -1)]
    # Seeded detected_at spans the last five days, all inside the window.
    open_rows = [v for v in db.query(Vulnerability).all()
                 if v.status == VulnerabilityStatus.open and v.asset_id != 1]
    assert sum(h["count"] for h in hist) == len(open_rows)
    print("✅ history is seven labelled days and totals the open findings")


def test_with_total_false_skips_the_count_only():
    """Same items, `total` null, one statement fewer."""
    engine, db = _db(hosts=20)
    # Called directly, so every Query() default has to be spelled out.
    kw = dict(severity=None, status="open", exploitable=None, kev_only=None,
              euvd_only=None, eu_critical=None, in_any_catalog=None,
              in_both_catalogs=None, epss_min=None, source=None,
              finding_type=None, cross_confirmed=None, asset_id=None, search=None, package=None,
              include_inactive_assets=False, distinct_cve=True,
              # detected_at, not published_date: that branch sorts with
              # split_part/regex, which is Postgres-only and never reaches the
              # count this test is about.
              sort_by="detected_at", sort_order="desc", limit=12, offset=0,
              db=db, current_user=_User())
    stmts = []
    event.listen(engine, "before_cursor_execute",
                 lambda conn, cur, stmt, *a: stmts.append(stmt))

    with_total = list_vulnerabilities(**kw)
    n_with = len(stmts)
    stmts.clear()
    without = list_vulnerabilities(with_total=False, **kw)
    n_without = len(stmts)

    def _ids(r):
        return [(v["cve_id"] if isinstance(v, dict) else v.cve_id) for v in r["items"]]
    assert _ids(with_total) == _ids(without)
    assert with_total["total"] > 0 and without["total"] is None
    assert n_without < n_with, f"count not skipped: {n_without} vs {n_with} statements"
    print(f"✅ same items, total skipped, {n_with} → {n_without} statements")


if __name__ == "__main__":
    test_statement_count_is_flat_and_small()
    test_numbers_match_the_per_count_semantics()
    test_seven_day_history_still_has_seven_days()
    test_with_total_false_skips_the_count_only()
    print("\nAll dashboard stats query tests passed.")