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 "DiscoverySeedContribution"')
        cur.execute('DELETE FROM "ScanWorkItem"')
        cur.execute('DELETE FROM "UnmatchedAlbum"')
        cur.execute('DELETE FROM "ApiCache"')
        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 "DiscoverySeedContribution"')
        cur.execute('DELETE FROM "ScanWorkItem"')
        cur.execute('DELETE FROM "UnmatchedAlbum"')
        cur.execute('DELETE FROM "ApiCache"')
        cur.execute('DELETE FROM "Config"')
    connection.commit()
    connection.close()


def insert_request(conn, artist="Artist", album="Album", force=False):
    """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, force, "createdAt") '
            "VALUES (gen_random_uuid()::text, %s, %s, 'pending', %s, now()) RETURNING id",
            (artist, album, force),
        )
        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