test_reader.py 3.45 KB
from unittest.mock import MagicMock

from etl_to_crawler.reader import fetch_platform_records, select_primary_record


def _make_cursor(rows):
    cur = MagicMock()
    cur.__enter__ = MagicMock(return_value=cur)
    cur.__exit__ = MagicMock(return_value=False)
    cur.fetchall.return_value = rows
    return cur


def test_fetch_platform_records_keeps_one_record_per_song_and_platform():
    conn = MagicMock()
    conn.cursor.return_value = _make_cursor([
        {
            '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': 101,
            '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',
        },
        {
            'source_song_id': 10,
            'record_id': 102,
            'platform': '4',
            'platform_unique_key': '300',
            'platform_mid': None,
            'album_audio_id': None,
            'is_main_version': 1,
            'is_high': 1,
            'pub_time': '2022-01-01',
        },
    ])

    result = fetch_platform_records(conn, [10])

    assert len(result) == 3
    assert {row['record_id'] for row in result} == {100, 101, 102}


def test_select_primary_record_prefers_main_version_then_is_high():
    rows = [
        {
            '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': 101,
            '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',
        },
        {
            'source_song_id': 10,
            'record_id': 102,
            'platform': '4',
            'platform_unique_key': '300',
            'platform_mid': None,
            'album_audio_id': None,
            'is_main_version': 1,
            'is_high': 1,
            'pub_time': '2022-01-01',
        },
    ]

    assert select_primary_record(rows)['record_id'] == 102


def test_select_primary_record_uses_earliest_pub_time_when_priority_ties():
    rows = [
        {
            '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': 101,
            'platform': '2',
            'platform_unique_key': '200',
            'platform_mid': 'kg-hash',
            'album_audio_id': None,
            'is_main_version': 0,
            'is_high': 0,
            'pub_time': '2019-01-01',
        },
    ]

    assert select_primary_record(rows)['record_id'] == 101