test_fix_record_titles.py 7.25 KB
from openpyxl import load_workbook
from datetime import datetime, timezone

from fix_record_titles import (
    InputRow,
    _archive_recording_sql,
    _archive_source_sql,
    _crawler_sql,
    _rollback_sql,
    _row_value,
    build_backup_records,
    build_report_rows,
    build_fix_rows,
    fetch_input_rows,
    load_backup,
    write_backup_atomic,
    write_report,
)


class Cursor:
    def __init__(self, rows):
        self.rows = rows
        self.calls = []

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return False

    def execute(self, sql, params=None):
        self.calls.append((sql, params))

    def fetchall(self):
        return self.rows


class Connection:
    def __init__(self, rows):
        self.raw_cursor = Cursor(rows)

    def cursor(self):
        return self.raw_cursor


def test_fetch_input_rows_reads_archive_push_relations_directly():
    connection = Connection([
        (1, "1", 101, "REC-1"),
        (2, "2", 102, "REC-2"),
    ])

    rows, stats = fetch_input_rows(connection)

    assert rows == [
        InputRow(1, "1", 101, "REC-1"),
        InputRow(2, "2", 102, "REC-2"),
    ]
    assert stats == {"total": 2, "missing_key": 0, "unsupported_platform": 0}
    sql = connection.raw_cursor.calls[0][0]
    assert "FROM yinyan_song_records" in sql
    assert "WHERE is_archive_push = TRUE" in sql


def test_fetch_input_rows_applies_limit():
    connection = Connection([(1, "1", 101, "REC-1")])

    fetch_input_rows(connection, limit=10)

    sql, params = connection.raw_cursor.calls[0]
    assert "LIMIT %s" in sql
    assert params == (10,)


def test_build_fix_rows_uses_spider_title_for_normalization():
    input_row = InputRow(1, "1", 101, "REC-1")

    rows, fix_stats = build_fix_rows(
        [input_row],
        {("1", 101): "七里香(Live版)"},
    )

    assert fix_stats["missing_spider"] == 0
    assert rows[0].spider_title == "七里香(Live版)"
    assert rows[0].title_norm == "七里香"
    assert rows[0].version == "live"
    assert _row_value(rows[0], "archive_platform") == "qqmusic"
    assert _row_value(rows[0], "archive_version") == "live"


def test_build_fix_rows_skips_missing_spider_title():
    input_row = InputRow(1, "4", 101, "REC-1")

    rows, fix_stats = build_fix_rows([input_row], {})

    assert rows == []
    assert fix_stats["missing_spider"] == 1


def test_build_fix_rows_temporarily_excludes_exact_spider_title():
    excluded = 'Waiting On A Wish (From "Disney\'s Snow White"/Soundtrack Version|Reprise)'
    input_rows = [
        InputRow(1, "1", 101, "REC-1"),
        InputRow(2, "1", 102, "REC-2"),
    ]
    spider_titles = {
        ("1", 101): excluded,
        ("1", 102): "Waiting On A Wish",
    }

    rows, fix_stats = build_fix_rows(
        input_rows,
        spider_titles,
        exclude_titles={excluded},
    )

    assert [row.platform_song_id for row in rows] == [102]
    assert fix_stats["excluded_title"] == 1


def test_update_sql_uses_spider_values_and_only_updates_differences():
    recording_sql, recording_params = _archive_recording_sql(1, True)
    source_sql, source_params = _archive_source_sql(1, True)
    crawler_sql, crawler_params = _crawler_sql("crawler_qqmusic_songs", 1, True)

    assert recording_params == ["recording_id", "spider_title", "title_norm", "archive_version"]
    assert "title_norm = source.new_title_norm" in recording_sql
    assert "IS DISTINCT FROM source.new_title" in recording_sql
    assert "target.platform_song_id = source.platform_song_id" in source_sql
    assert "spider_title" in source_params
    assert crawler_params == ["platform_song_id", "spider_title", "version"]
    assert "source.platform_song_id::bigint" in crawler_sql
    assert "version = source.new_version" in crawler_sql
    assert "RETURNING target.recording_id" in recording_sql
    assert "RETURNING target.platform_song_id" in crawler_sql
    locked_sql, _ = _archive_recording_sql(1, False, lock=True)
    assert "FOR UPDATE OF target" in locked_sql


def test_report_lists_only_returned_target_keys(tmp_path):
    fix = build_fix_rows(
        [InputRow(1, "1", 101, "REC-1")],
        {("1", 101): "七里香(Live版)"},
    )[0][0]
    matched = {
        "archive_recording": [
            ("REC-1", "旧标题", "七里香(Live版)", "旧标题", "七里香", None, "live"),
        ],
        "archive_recording_platform_source": [
            ("REC-1", "qqmusic", "101", "旧标题", "七里香(Live版)", None, "七里香", None, "live"),
        ],
        "crawler_qqmusic_songs": [
            (101, "旧标题", "七里香(Live版)", "", "live"),
        ],
        "crawler_kugou_songs": [],
        "crawler_netease_songs": [],
    }

    rows = build_report_rows([fix], matched, apply=False)
    output = tmp_path / "report.xlsx"
    write_report(output, rows)

    worksheet = load_workbook(output).active
    assert len(rows) == 3
    assert worksheet.max_row == 4
    assert worksheet["A1"].value == "database"
    assert worksheet["H1"].value == "old_title"
    assert worksheet["I1"].value == "new_title"
    assert worksheet["H2"].value == "旧标题"
    assert worksheet["I2"].value == "七里香(Live版)"
    assert worksheet["J3"].value == "<NULL>"
    assert worksheet["J3"].fill.fgColor.rgb == "00FFC7CE"
    assert worksheet["K3"].fill.fgColor.rgb == "00C6EFCE"
    assert worksheet["N2"].value == "would_update"


def test_backup_round_trip_preserves_null_empty_and_updated_at(tmp_path):
    fix = build_fix_rows(
        [InputRow(1, "1", 101, "REC-1")],
        {("1", 101): "七里香(Live版)"},
    )[0][0]
    updated_at = datetime(2026, 7, 15, 10, 0, tzinfo=timezone.utc)
    preview = {
        "archive_recording": [
            ("REC-1", "旧标题", "七里香(Live版)", None, "七里香", None, "live", updated_at),
        ],
        "archive_recording_platform_source": [],
        "crawler_qqmusic_songs": [
            (101, "旧标题", "七里香(Live版)", "", "live", updated_at.replace(tzinfo=None)),
        ],
        "crawler_kugou_songs": [],
        "crawler_netease_songs": [],
    }
    records = build_backup_records([fix], preview, "run-1")
    path = tmp_path / "backup.jsonl"

    write_backup_atomic(path, records, "run-1")
    metadata, loaded = load_backup(path)

    assert metadata["change_count"] == 2
    assert loaded[0]["old"]["title_norm"] is None
    assert loaded[1]["old"]["version"] == ""
    assert loaded[0]["old"]["updated_at"] == updated_at.isoformat()

    try:
        write_backup_atomic(path, records, "run-2")
    except FileExistsError:
        pass
    else:
        raise AssertionError("已有备份不应被覆盖")


def test_rollback_sql_restores_old_values_with_new_value_guard():
    archive_sql, _ = _rollback_sql("archive_recording", 1)
    crawler_sql, _ = _rollback_sql("crawler_qqmusic_songs", 1)

    assert "SET title = source.old_title" in archive_sql
    assert "updated_at = source.old_updated_at::timestamptz" in archive_sql
    assert "target.title IS NOT DISTINCT FROM source.new_title" in archive_sql
    assert "target.title_norm IS NOT DISTINCT FROM source.new_title_norm" in archive_sql
    assert "updated_at = source.old_updated_at::timestamp" in crawler_sql
    assert "target.version IS NOT DISTINCT FROM source.new_version" in crawler_sql