test_runner.py 10.6 KB
from unittest.mock import MagicMock

from etl_to_crawler import runner


class _Cursor:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False

    def execute(self, *args, **kwargs):
        return None


class _PgConnection:
    def __init__(self):
        self.cur = _Cursor()
        self.commits = 0
        self.closed = False

    def cursor(self):
        return self.cur

    def commit(self):
        self.commits += 1

    def close(self):
        self.closed = True


class _Connection:
    def close(self):
        return None


def test_run_imports_all_platform_records_but_writes_one_primary_yinyan_relation(monkeypatch):
    pg_conn = _PgConnection()
    processors = {
        '1': MagicMock(return_value={'platform': 'qq', 'platform_song_id': 100, 'mid': 'qq-mid', 'title': '歌'}),
        '2': MagicMock(return_value={'platform': 'kugou', 'platform_song_id': 200, 'hash': 'kg-hash', 'title': '歌'}),
    }
    yinyan_writer = MagicMock()

    monkeypatch.setattr(runner, 'get_hk_songs_conn', lambda: _Connection())
    monkeypatch.setattr(runner, 'get_source_conn', lambda: _Connection())
    monkeypatch.setattr(runner, 'get_spider_conn', lambda: _Connection())
    monkeypatch.setattr(runner, 'get_pg_conn', lambda: pg_conn)
    monkeypatch.setattr(runner, 'get_oss_bucket', lambda: object())
    monkeypatch.setattr(runner, 'fetch_pending_yinyan_song_records', lambda cur, batch_size: [
        {'song_id': 10, 'record_id': 200},
    ] if pg_conn.commits == 0 else [])
    monkeypatch.setattr(runner, 'fetch_hk_songs_by_source_ids', lambda conn, song_ids: {10: {
        'source_song_id': 10,
        'name': '歌',
        'audio_url': 'https://example.com/a.mp3',
        'singer': '歌手',
    }})
    platform_records = [
        {
            'source_song_id': 10,
            'record_id': 100,
            'platform': '1',
            'platform_unique_key': 'qq-mid',
            'platform_mid': '100',
            'album_audio_id': None,
            'is_main_version': 0,
            'is_high': 1,
            'pub_time': '2020-01-01',
        },
        {
            'source_song_id': 10,
            'record_id': 200,
            'platform': '2',
            'platform_unique_key': '200',
            'platform_mid': 'kg-hash',
            'album_audio_id': None,
            'is_main_version': 1,
            'is_high': 0,
            'pub_time': '2021-01-01',
        },
    ]
    monkeypatch.setattr(runner, 'fetch_platform_records', lambda conn, song_ids: platform_records)
    monkeypatch.setattr(runner, 'fetch_all_platform_records', lambda conn, song_ids: platform_records)
    monkeypatch.setattr(runner, '_pick_record_with_singer', lambda records, spider_conn: records[0])
    monkeypatch.setattr(runner, '_PROCESSORS', processors)
    monkeypatch.setattr(runner, 'upsert_yinyan_song_records', yinyan_writer)

    runner.run(['1', '2'])

    processors['1'].assert_called_once()
    processors['2'].assert_called_once()
    yinyan_writer.assert_called_once_with(pg_conn.cur, [{
        'song_id': 10,
        'record_id': 100,
        'platform': '1',
        'platform_song_id': 100,
    }])
    assert pg_conn.commits == 1


def test_initialize_yinyan_song_records_inserts_primary_records(monkeypatch):
    pg_conn = _PgConnection()
    inserted = []

    monkeypatch.setattr(runner, 'get_hk_songs_conn', lambda: _Connection())
    monkeypatch.setattr(runner, 'get_source_conn', lambda: _Connection())
    monkeypatch.setattr(runner, 'get_pg_conn', lambda: pg_conn)
    monkeypatch.setattr(runner, 'fetch_existing_yinyan_song_ids', lambda cur: set())

    monkeypatch.setattr(runner, 'iter_hk_songs_batches', lambda conn, batch_size: [[
        {'id': 50, 'source_song_id': 10, 'name': '歌', 'audio_url': 'https://example.com/a.mp3', 'singer': '歌手'},
    ]])
    monkeypatch.setattr(runner, 'fetch_platform_records', lambda conn, song_ids: [
        {
            'source_song_id': 10,
            'record_id': 100,
            'platform': '1',
            'platform_unique_key': 'qq-mid',
            'platform_mid': '100',
            'album_audio_id': None,
            'is_main_version': 0,
            'is_high': 0,
            'pub_time': '2020-01-01',
        },
        {
            'source_song_id': 10,
            'record_id': 200,
            'platform': '2',
            'platform_unique_key': '200',
            'platform_mid': 'kg-hash',
            'album_audio_id': None,
            'is_main_version': 1,
            'is_high': 0,
            'pub_time': '2021-01-01',
        },
    ])
    monkeypatch.setattr(runner, 'insert_yinyan_song_records', lambda cur, rows: inserted.extend(rows))

    runner.initialize_yinyan_song_records(['1', '2'])

    assert inserted == [{'song_id': 10, 'record_id': 200, 'platform': '2'}]
    assert pg_conn.commits == 1


def test_backfill_yinyan_record_platforms_updates_missing_platform_rows(monkeypatch):
    pg_conn = _PgConnection()
    updated = []

    monkeypatch.setattr(runner, 'get_source_conn', lambda: _Connection())
    monkeypatch.setattr(runner, 'get_pg_conn', lambda: pg_conn)
    monkeypatch.setattr(runner, 'fetch_yinyan_records_missing_platform', lambda cur, batch_size: [
        {'song_id': 10, 'record_id': 100},
        {'song_id': 11, 'record_id': 101},
    ] if pg_conn.commits == 0 else [])
    monkeypatch.setattr(runner, 'fetch_record_platforms', lambda conn, record_ids: {
        100: '1',
        101: '2',
    })
    monkeypatch.setattr(runner, 'update_yinyan_record_platforms', lambda cur, rows: updated.extend(rows))

    runner.backfill_yinyan_record_platforms()

    assert updated == [
        {'song_id': 10, 'record_id': 100, 'platform': '1'},
        {'song_id': 11, 'record_id': 101, 'platform': '2'},
    ]
    assert pg_conn.commits == 1


def test_process_netease_builds_album_json_for_song_insert(monkeypatch):
    pg_cur = MagicMock()
    pg_cur.fetchone.return_value = ('song-uuid',)
    inserted_songs = []

    monkeypatch.setattr(runner, 'fetch_netease_songs', lambda conn, song_ids: {
        300: {
            'id': 300,
            'album_id': 20,
            'album_cover': 'https://example.com/album.jpg',
            'album_title': '专辑',
            'album_intro': '简介',
            'album_type': '专辑类型',
            'company_id': 7,
            'company': '唱片公司',
            'is_owner': 1,
            'album_published_at': '2020-01-01',
            'cover': 'https://example.com/cover.jpg',
            'title': '化风行万里 (DJ默涵版)',
            'duration': 180,
            'lyric': '[00:01.00]歌词',
            'composer_name': '曲作者',
            'lyricist_name': '词作者',
            'platform_index_url': None,
            'published_at': '2020-01-02',
        },
    })
    monkeypatch.setattr(runner, 'fetch_netease_singers', lambda conn, song_ids: {})
    monkeypatch.setattr(runner, '_safe_transfer', lambda url, oss_key, bucket, base_url: url)
    monkeypatch.setattr(runner, '_safe_transfer_audio', lambda url, oss_key, bucket, base_url: (url, 'audio-md5'))
    monkeypatch.setattr(runner, 'upsert_netease_singers', lambda cur, singers: None)
    monkeypatch.setattr(runner, 'upsert_netease_albums', lambda cur, albums: None)
    monkeypatch.setattr(runner, 'upsert_netease_songs', lambda cur, songs: inserted_songs.extend(songs))
    monkeypatch.setattr(runner, 'upsert_netease_singer_songs', lambda cur, pairs: None)
    monkeypatch.setattr(runner, 'upsert_netease_singer_albums', lambda cur, pairs: None)

    runner._process_netease(
        {
            'name': '词曲名',
            'audio_url': 'https://example.com/audio.mp3',
            'lyrics_url': 'https://example.com/lyric.lrc',
            'cover_url': '',
            'composer': '词曲曲作者',
            'lyricist': '词曲词作者',
            'issue_time': '2019-01-01',
            'song_time': 120,
        },
        {'platform_unique_key': '300'},
        spider_conn=object(),
        pg_cur=pg_cur,
        bucket=object(),
        base_url='https://bucket.example.com',
    )

    assert inserted_songs[0]['album_json']
    assert '"id": 20' in inserted_songs[0]['album_json']
    assert '"title": "专辑"' in inserted_songs[0]['album_json']
    assert inserted_songs[0]['title'] == '化风行万里'
    assert inserted_songs[0]['version'] == 'DJ默涵版'


def test_process_netease_keeps_timestamped_lyric_and_uploads_plain_lyric(monkeypatch):
    pg_cur = MagicMock()
    pg_cur.fetchone.return_value = ('song-uuid',)
    inserted_songs = []
    uploaded = {}

    class Bucket:
        def put_object(self, key, body, headers=None):
            uploaded['key'] = key
            uploaded['body'] = body

    monkeypatch.setattr(runner, 'fetch_netease_songs', lambda conn, song_ids: {
        300: {
            'id': 300,
            'album_id': None,
            'cover': 'https://example.com/cover.jpg',
            'title': '录音标题',
            'duration': 180,
            'lyric': '[ti:歌名]\n[00:01.00]第一句\n[00:02.00]第二句',
            'composer_name': '曲作者',
            'lyricist_name': '词作者',
            'platform_index_url': None,
            'published_at': '2020-01-02',
        },
    })
    monkeypatch.setattr(runner, 'fetch_netease_singers', lambda conn, song_ids: {})
    monkeypatch.setattr(runner, '_safe_transfer', lambda url, oss_key, bucket, base_url: url)
    monkeypatch.setattr(runner, '_safe_transfer_audio', lambda url, oss_key, bucket, base_url: (url, 'audio-md5'))
    monkeypatch.setattr(runner, 'upsert_netease_singers', lambda cur, singers: None)
    monkeypatch.setattr(runner, 'upsert_netease_albums', lambda cur, albums: None)
    monkeypatch.setattr(runner, 'upsert_netease_songs', lambda cur, songs: inserted_songs.extend(songs))
    monkeypatch.setattr(runner, 'upsert_netease_singer_songs', lambda cur, pairs: None)

    runner._process_netease(
        {
            'name': '词曲名',
            'audio_url': 'https://example.com/audio.mp3',
            'lyrics_url': 'https://example.com/original.lrc',
            'cover_url': '',
            'composer': '词曲曲作者',
            'lyricist': '词曲词作者',
            'issue_time': '2019-01-01',
            'song_time': 120,
        },
        {'platform_unique_key': '300'},
        spider_conn=object(),
        pg_cur=pg_cur,
        bucket=Bucket(),
        base_url='https://bucket.example.com',
    )

    assert inserted_songs[0]['lyric'] == '[ti:歌名]\n[00:01.00]第一句\n[00:02.00]第二句'
    assert inserted_songs[0]['audio_md5'] == 'audio-md5'
    assert inserted_songs[0]['lyric_url'] == 'https://bucket.example.com/crawler/netease/lyric/300.txt'
    assert uploaded['body'].decode('utf-8') == '第一句\n第二句'