Lyra/worker/tests/conftest.py

import os

os.environ.setdefault("DATABASE_URL", "postgresql://lyra:lyra@localhost:5432/lyra_test")

import psycopg
import pytest


def _require_test_db(dsn: str) -> None:
    name = dsn.rsplit("/", 1)[-1].split("?")[0]
    if not name.endswith("_test"):
        raise RuntimeError(
            f"refusing to run destructive test fixtures against non-test database '{name}'. "
            "Point DATABASE_URL at a database whose name ends in '_test' (e.g. lyra_test)."
        )


@pytest.fixture()
def conn():
    dsn = os.environ["DATABASE_URL"]
    _require_test_db(dsn)
    connection = psycopg.connect(dsn)
    # Clean up any pre-existing rows before the test runs (Job cascades from Request).
    with connection.cursor() as cur:
        cur.execute('DELETE FROM "Job"')
        cur.execute('DELETE FROM "Request"')
        cur.execute('DELETE FROM "LibraryItem"')
        cur.execute('DELETE FROM "MonitoredRelease"')
        cur.execute('DELETE FROM "WatchedArtist"')
        cur.execute('DELETE FROM "DiscoverySuggestion"')
        cur.execute('DELETE FROM "Config"')
    connection.commit()
    yield connection
    # A test may have left the transaction aborted (e.g. asserting on a
    # UniqueViolation via pytest.raises without rolling back); reset it so
    # the cleanup below can run.
    connection.rollback()
    # Clean up rows created during the test (Job cascades from Request).
    with connection.cursor() as cur:
        cur.execute('DELETE FROM "Job"')
        cur.execute('DELETE FROM "Request"')
        cur.execute('DELETE FROM "LibraryItem"')
        cur.execute('DELETE FROM "MonitoredRelease"')
        cur.execute('DELETE FROM "WatchedArtist"')
        cur.execute('DELETE FROM "DiscoverySuggestion"')
        cur.execute('DELETE FROM "Config"')
    connection.commit()
    connection.close()


def insert_request(conn, artist="Artist", album="Album"):
    """Insert a Request + its Job (state 'requested') and return the job id."""
    with conn.cursor() as cur:
        cur.execute(
            'INSERT INTO "Request" (id, artist, album, status, "createdAt") '
            "VALUES (gen_random_uuid()::text, %s, %s, 'pending', now()) RETURNING id",
            (artist, album),
        )
        request_id = cur.fetchone()[0]
        cur.execute(
            'INSERT INTO "Job" (id, "requestId", state, "currentStage", attempts, "createdAt", "updatedAt") '
            "VALUES (gen_random_uuid()::text, %s, 'requested', 'intake', 0, now(), now()) RETURNING id",
            (request_id,),
        )
        job_id = cur.fetchone()[0]
    conn.commit()
    return job_id


def insert_watched_artist(conn, mbid="mbid", name="Artist", auto_monitor_future=False,
                          monitor_from_sql="now()", last_polled_sql="NULL",
                          last_discovered_sql="NULL"):
    """Insert a WatchedArtist; *_sql args are raw SQL time expressions."""
    with conn.cursor() as cur:
        cur.execute(
            'INSERT INTO "WatchedArtist" (id, mbid, name, "autoMonitorFuture", '
            f'"monitorFrom", "lastPolledAt", "lastDiscoveredAt", "createdAt") '
            f"VALUES (gen_random_uuid()::text, %s, %s, %s, {monitor_from_sql}, "
            f"{last_polled_sql}, {last_discovered_sql}, now()) RETURNING id",
            (mbid, name, auto_monitor_future),
        )
        artist_id = cur.fetchone()[0]
    conn.commit()
    return artist_id


def insert_monitored_release(conn, watched_artist_id=None, artist_mbid="mbid", artist_name="Artist",
                             rg_mbid="rg", album="Album", primary_type=None, secondary_types=None,
                             first_release_date=None, monitored=False, state="wanted",
                             current_quality_class=None, first_grabbed_sql="NULL",
                             last_searched_sql="NULL"):
    """Insert a MonitoredRelease; *_sql args are raw SQL time expressions."""
    with conn.cursor() as cur:
        cur.execute(
            'INSERT INTO "MonitoredRelease" (id, "watchedArtistId", "artistMbid", "artistName", '
            '"rgMbid", album, "primaryType", "secondaryTypes", "firstReleaseDate", monitored, '
            f'state, "currentQualityClass", "firstGrabbedAt", "lastSearchedAt", "createdAt") '
            f"VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, "
            f"{first_grabbed_sql}, {last_searched_sql}, now()) RETURNING id",
            (watched_artist_id, artist_mbid, artist_name, rg_mbid, album, primary_type,
             secondary_types if secondary_types is not None else [], first_release_date,
             monitored, state, current_quality_class),
        )
        rel_id = cur.fetchone()[0]
    conn.commit()
    return rel_id


def insert_discovery_suggestion(conn, kind="artist", artist_mbid="cand", artist_name="Cand",
                                rg_mbid=None, album=None, score=1.0, seed_count=1,
                                sources=None, status="pending"):
    """Insert a DiscoverySuggestion; returns its id."""
    dedupe = f"{kind}:{artist_mbid}:{rg_mbid or ''}"
    with conn.cursor() as cur:
        cur.execute(
            'INSERT INTO "DiscoverySuggestion" (id, kind, "artistMbid", "artistName", '
            '"rgMbid", album, score, "seedCount", sources, status, "dedupeKey", '
            '"createdAt", "updatedAt") '
            'VALUES (gen_random_uuid()::text, %s::"DiscoveryKind", %s, %s, %s, %s, %s, %s, %s, '
            '%s::"SuggestionStatus", %s, now(), now()) RETURNING id',
            (kind, artist_mbid, artist_name, rg_mbid, album, score, seed_count,
             sources if sources is not None else ["listenbrainz"], status, dedupe),
        )
        sid = cur.fetchone()[0]
    conn.commit()
    return sid