test_oss.py 3.54 KB
from unittest.mock import MagicMock, patch
import math
from etl_to_crawler.oss import transfer_url, transfer_url_with_md5

ARCHIVE_URL = "https://archive-dev.oss-cn-beijing.aliyuncs.com/some/path.mp3"
OTHER_URL = "https://hikoon-data-platform.oss-cn-beijing.aliyuncs.com/qq-audio/abc.mp3"
BASE_URL = "https://archive-dev.oss-cn-beijing.aliyuncs.com"

def test_already_archive_dev_returns_as_is():
    bucket = MagicMock()
    result = transfer_url(ARCHIVE_URL, "any/key.mp3", bucket, BASE_URL)
    assert result == ARCHIVE_URL
    bucket.put_object.assert_not_called()

def test_empty_url_returns_empty():
    bucket = MagicMock()
    result = transfer_url('', "any/key.mp3", bucket, BASE_URL)
    assert result == ''
    bucket.put_object.assert_not_called()

def test_none_url_returns_empty():
    bucket = MagicMock()
    result = transfer_url(None, "any/key.mp3", bucket, BASE_URL)
    assert result == ''


def test_invalid_nan_url_returns_empty_without_download():
    bucket = MagicMock()
    with patch('etl_to_crawler.oss._http_get') as mock_get:
        assert transfer_url('nan', "any/key.mp3", bucket, BASE_URL) == ''
        assert transfer_url(math.nan, "any/key.mp3", bucket, BASE_URL) == ''
        assert transfer_url('not-a-url', "any/key.mp3", bucket, BASE_URL) == ''
        assert transfer_url_with_md5('nan', "any/key.mp3", bucket, BASE_URL) == ('', '')
    mock_get.assert_not_called()
    bucket.put_object.assert_not_called()


def test_external_url_downloads_and_uploads():
    bucket = MagicMock()
    fake_content = b"audio_bytes"
    with patch('etl_to_crawler.oss._http_get') as mock_get:
        mock_get.return_value.content = fake_content
        mock_get.return_value.raise_for_status = MagicMock()
        result = transfer_url(OTHER_URL, "crawler/qq/audio/abc.mp3", bucket, BASE_URL)
    bucket.put_object.assert_called_once_with("crawler/qq/audio/abc.mp3", fake_content)
    assert result == f"{BASE_URL}/crawler/qq/audio/abc.mp3"


def test_external_url_can_rewrite_download_base_to_internal_endpoint():
    bucket = MagicMock()
    public_source = "https://source-bucket.oss-cn-hangzhou.aliyuncs.com/audio/abc.mp3"
    internal_source_base = "https://source-bucket.oss-cn-hangzhou-internal.aliyuncs.com"
    fake_content = b"audio_bytes"
    with patch.dict('etl_to_crawler.oss.OSS_CONFIG', {
        'download_rewrite_from_base_url': 'https://source-bucket.oss-cn-hangzhou.aliyuncs.com',
        'download_base_url': internal_source_base,
    }):
        with patch('etl_to_crawler.oss._http_get') as mock_get:
            mock_get.return_value.content = fake_content
            mock_get.return_value.raise_for_status = MagicMock()
            result = transfer_url(public_source, "crawler/qq/audio/abc.mp3", bucket, BASE_URL)

    mock_get.assert_called_once_with(f"{internal_source_base}/audio/abc.mp3", timeout=30)
    bucket.put_object.assert_called_once_with("crawler/qq/audio/abc.mp3", fake_content)
    assert result == f"{BASE_URL}/crawler/qq/audio/abc.mp3"


def test_transfer_url_with_md5_hashes_downloaded_content_before_upload():
    bucket = MagicMock()
    fake_content = b"audio_bytes"
    with patch('etl_to_crawler.oss._http_get') as mock_get:
        mock_get.return_value.content = fake_content
        mock_get.return_value.raise_for_status = MagicMock()
        result, audio_md5 = transfer_url_with_md5(OTHER_URL, "crawler/qq/audio/abc.mp3", bucket, BASE_URL)
    bucket.put_object.assert_called_once_with("crawler/qq/audio/abc.mp3", fake_content)
    assert result == f"{BASE_URL}/crawler/qq/audio/abc.mp3"
    assert audio_md5 == "04d43544b267629d9089eaed3b847a99"