PacketStation/tests/server/notifyhelpers.py

"""Helpers for the notify tests: a module-scoped database, a small directory fake with the notify calls, users."""

from __future__ import annotations

import secrets
import uuid
from collections.abc import Iterator
from pathlib import Path
from typing import Any

from fakedir import FakeDirectory
from pshelpers import PG_URL, app_conn, make_settings

from packetstation.auth.ldap import AdUser
from packetstation.notify.routing import cn_of
from packetstation.settings import Settings

BASE_DN = "DC=pstest,DC=local"
GROUP_BASE_DN = f"OU=PacketStation,OU=Groups,OU=PESCO_FARMINGTON,{BASE_DN}"
OTHER_BASE_DN = f"OU=Groups,{BASE_DN}"
OUTSIDE = "Ops-Mail"


def fresh_db(tmp: Path, **over: Any) -> Iterator[Settings]:
    """A fresh migrated database (like conftest.pg_settings), dropped afterwards."""
    import psycopg
    from psycopg import sql
    from psycopg.conninfo import conninfo_to_dict

    from packetstation import db, migrate

    assert PG_URL
    info = conninfo_to_dict(PG_URL)
    mig_pw, app_pw = secrets.token_urlsafe(18), secrets.token_urlsafe(18)
    (tmp / "mig").write_text(mig_pw + "\n")
    (tmp / "app").write_text(app_pw + "\n")
    name = f"ps_test_{secrets.token_hex(4)}"
    s = make_settings(
        tmp,
        pg_host=info.get("host", "localhost"),
        pg_port=int(info.get("port", 5432)),
        pg_database=name,
        pg_migrate_password_file=tmp / "mig",
        pg_app_password_file=tmp / "app",
        pg_superuser_url=PG_URL,
        ldap_base_dn=BASE_DN,
        ldap_group_base_dn=GROUP_BASE_DN,
        smtp_from="[email protected]",
        **over,
    )
    migrate.bootstrap_roles(s)
    with psycopg.connect(PG_URL, autocommit=True) as su:
        for user, pw in (("ps_migrate", mig_pw), ("ps_app", app_pw)):
            su.execute(sql.SQL("ALTER ROLE {} LOGIN PASSWORD {}").format(sql.Identifier(user), sql.Literal(pw)))
    with db.connect_as("migrate", s, autocommit=True) as conn:
        migrate.run(conn, db.MIGRATIONS_DIR)
    try:
        yield s
    finally:
        db.close_pool()
        with psycopg.connect(PG_URL, autocommit=True) as su:
            su.execute(sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(sql.Identifier(name)))


class NotifyDir(FakeDirectory):
    """FakeDirectory plus members_of_dn and search_groups, by DN under GROUP_BASE_DN or OTHER_BASE_DN.

    OUTSIDE is a group outside the OU whose only member is qc.quinn."""

    def __init__(self) -> None:
        super().__init__()
        self.outside: dict[str, list[str]] = {OUTSIDE: ["qc.quinn"]}
        self.member_calls: list[str] = []

    @staticmethod
    def dn(name: str) -> str:
        return f"CN={name},{OTHER_BASE_DN if name == OUTSIDE else GROUP_BASE_DN}"

    def members_of_dn(self, group_dn: str) -> list[AdUser]:
        self._check()
        self.member_calls.append(group_dn)
        name = cn_of(group_dn)
        if name in self.outside:
            return [self._ad(self.user(sam)) for sam in self.outside[name]]
        if name.lower() not in {g.lower() for g in self.ou_groups}:
            return []
        return [self._ad(u) for u in self.users if name.lower() in self._closure(u)]

    def search_groups(self, q: str, limit: int = 25) -> list[dict[str, Any]]:
        self._check()
        key = q.strip().lower()
        names = sorted({*self.ou_groups, OUTSIDE}, key=str.lower)
        return [{"dn": self.dn(n), "name": n, "description": f"group {n}"} for n in names if key in n.lower()][:limit]


def add_users(settings: Settings, directory: FakeDirectory, sams: tuple[str, ...] | None = None) -> dict[str, str]:
    """users rows for the fake people (id = their GUID, email [email protected]). Returns sam -> id."""
    out: dict[str, str] = {}
    with app_conn(settings) as conn:
        for u in directory.users:
            if sams is not None and u.sam not in sams:
                continue
            conn.execute(
                "INSERT INTO users (id, sam_account, display_name, emp_no, email, active) VALUES (%s, %s, %s, %s, %s,"
                " %s) ON CONFLICT (id) DO NOTHING",
                (u.guid, u.sam, u.display, u.emp_no, f"{u.sam}@example.test", not u.disabled),
            )
            out[u.sam] = str(u.guid)
    return out


def clean(settings: Settings) -> int:
    """Empty the notification tables a test touches (notifications_sent is append-only: returns its max id)."""
    with app_conn(settings) as conn:
        conn.execute("DELETE FROM notification_outbox")
        conn.execute("DELETE FROM notification_events")
        conn.execute("DELETE FROM notification_exclusions")
        conn.execute("DELETE FROM notification_recipients")
        conn.execute("DELETE FROM notification_rules")
        conn.execute("DELETE FROM settings WHERE key = 'notify.last_build'")
        conn.execute("DELETE FROM drawing_lookups")
        return int(conn.execute("SELECT coalesce(max(id), 0) FROM notifications_sent").fetchone()[0])


def sent_since(settings: Settings, marker: int) -> list[dict[str, Any]]:
    from psycopg.rows import dict_row

    with app_conn(settings) as conn, conn.cursor(row_factory=dict_row) as cur:
        return [dict(r) for r in cur.execute("SELECT * FROM notifications_sent WHERE id > %s ORDER BY id", (marker,))]


def new_id() -> uuid.UUID:
    return uuid.uuid4()