import io import math import os import re import tempfile from typing import Any from unittest.mock import Mock, call, patch import httpx import pytest from fsspec import AbstractFileSystem from fsspec.implementations.http import HTTPFileSystem from owsclient import ImpersonationOwsClient, OwsClient from owsclient.test import OwsClientMock from s3fs import S3FileSystem from ows_assets_uploader import uploader from ows_assets_uploader.uploader import AssetsUploader, _OwsClientWrapper @pytest.fixture def ows_client(ows_client_mock: OwsClientMock) -> OwsClient: return OwsClient( environment="test", service_name="python-ows-assets-upload", ) @pytest.fixture def assets_uploader(ows_client: OwsClient) -> AssetsUploader: return AssetsUploader(ows_client=ows_client) @pytest.fixture def impersonation_ows_client(ows_client_mock: OwsClientMock) -> ImpersonationOwsClient: return ImpersonationOwsClient( environment="test", service_name="python-ows-assets-upload", m2m_token_manager=ows_client_mock.mock_impersonation_m2m_token_manager, ) @pytest.fixture def impersonated_identity_uuid() -> str: return "test-identity-uuid" @pytest.fixture def obo_assets_uploader( impersonation_ows_client: ImpersonationOwsClient, impersonated_identity_uuid: str, ) -> AssetsUploader: return AssetsUploader( ows_client=impersonation_ows_client, impersonated_identity_uuid=impersonated_identity_uuid, ) class TestEnvironmentFunctions: def setup_method(self) -> None: uploader._env = None def teardown_method(self) -> None: uploader._env = None def test_set_env_after_already_set(self) -> None: uploader.set_env("qa") with pytest.raises( RuntimeError, match=re.escape("Environment already set to qa. It can only be set once."), ): uploader.set_env("prod") assert uploader.get_env() == "qa" @pytest.mark.parametrize("valid_env", uploader._VALID_ENVIRONMENTS) def test_set_env_valid_environments(self, valid_env: str) -> None: uploader.set_env(valid_env) assert uploader.get_env() == valid_env def test_set_env_invalid_environment(self) -> None: with pytest.raises( ValueError, match=re.escape("Invalid environment: invalid. Must be one of") ): uploader.set_env("invalid") def test_get_env_not_set(self) -> None: with pytest.raises( RuntimeError, match=re.escape("Environment not set. Call set_env() to set it."), ): uploader.get_env() class TestSetM2MTokenManager: def setup_method(self) -> None: uploader._m2m_token_manager = None def teardown_method(self) -> None: uploader._m2m_token_manager = None def test_set_m2m_token_manager(self) -> None: m2m_token_manager = Mock() uploader.set_m2m_token_manager(m2m_token_manager) assert uploader._m2m_token_manager is m2m_token_manager def test_set_m2m_token_manager_after_already_set(self) -> None: uploader.set_m2m_token_manager(Mock()) with pytest.raises( RuntimeError, match=re.escape("M2M token manager already set. It can only be set once."), ): uploader.set_m2m_token_manager(Mock()) class TestBatched: def test_batched_empty(self) -> None: result: list[tuple[int, ...]] = list(uploader._batched([], 3)) assert result == [] def test_batched_single_batch(self) -> None: result = list(uploader._batched([1, 2, 3], 5)) assert result == [(1, 2, 3)] def test_batched_multiple_batches(self) -> None: result = list(uploader._batched([1, 2, 3, 4, 5, 6, 7], 3)) assert result == [(1, 2, 3), (4, 5, 6), (7,)] def test_batched_exact_multiple(self) -> None: result = list(uploader._batched([1, 2, 3, 4, 5, 6], 3)) assert result == [(1, 2, 3), (4, 5, 6)] class TestTruncateMiddle: @pytest.mark.parametrize( ("test_description", "string", "max_length", "expected"), [ ("under even limit", "abcde", 6, "abcde"), ("at even limit", "abcdef", 6, "abcdef"), ("over even limit", "abcdefg", 6, "a...fg"), ("under odd limit", "abcdef", 7, "abcdef"), ("at odd limit", "abcdefg", 7, "abcdefg"), ("over odd limit", "abcdefgh", 7, "ab...gh"), ], ) def test_truncate_middle( self, test_description: str, string: str, max_length: int, expected: str ) -> None: result = uploader._truncate_middle(string, max_length) assert result == expected class TestGetOwsClient: def setup_method(self) -> None: uploader._env = None uploader._m2m_token_manager = None uploader._get_ows_client.cache_clear() def teardown_method(self) -> None: uploader._env = None uploader._m2m_token_manager = None uploader._get_ows_client.cache_clear() @patch("ows_assets_uploader.uploader.OwsClient") def test_get_ows_client(self, mock_ows_client: Mock) -> None: uploader.set_env("qa") client1 = uploader._get_ows_client() mock_ows_client.assert_called_once_with( environment="qa", service_name="python-ows-assets-upload", retries=5, m2m_token_manager=None, ) client2 = uploader._get_ows_client() assert client1 is client2 @patch("ows_assets_uploader.uploader.OwsClient") def test_get_ows_client_uses_registered_m2m_token_manager( self, mock_ows_client: Mock ) -> None: uploader.set_env("qa") m2m_token_manager = Mock() uploader.set_m2m_token_manager(m2m_token_manager) uploader._get_ows_client() mock_ows_client.assert_called_once_with( environment="qa", service_name="python-ows-assets-upload", retries=5, m2m_token_manager=m2m_token_manager, ) class TestGetHttpxClient: def teardown_method(self) -> None: uploader._get_httpx_client.cache_clear() def test_get_httpx_client_returns_singleton(self) -> None: client1 = uploader._get_httpx_client() client2 = uploader._get_httpx_client() assert client1 is client2 assert isinstance(client1, httpx.Client) class TestAssetsUploaderInit: def test_init_creates_httpx_client(self, ows_client: OwsClient) -> None: assets_uploader = AssetsUploader(ows_client=ows_client) assert isinstance(assets_uploader._httpx_client, httpx.Client) assert isinstance(assets_uploader._ows_client, _OwsClientWrapper) def test_init_with_impersonation_client( self, impersonation_ows_client: ImpersonationOwsClient, impersonated_identity_uuid: str, ) -> None: assets_uploader = AssetsUploader( ows_client=impersonation_ows_client, impersonated_identity_uuid=impersonated_identity_uuid, ) assert isinstance(assets_uploader._httpx_client, httpx.Client) assert isinstance(assets_uploader._ows_client, _OwsClientWrapper) def test_init_impersonation_client_requires_identity_uuid( self, impersonation_ows_client: ImpersonationOwsClient ) -> None: with pytest.raises( ValueError, match=re.escape( "impersonated_identity_uuid is required when using ImpersonationOwsClient" ), ): AssetsUploader(ows_client=impersonation_ows_client) def test_init_standard_client_rejects_impersonated_identity_uuid( self, ows_client: OwsClient ) -> None: with pytest.raises( ValueError, match=re.escape( "impersonated_identity_uuid should not be provided when using standard OwsClient" ), ): AssetsUploader( ows_client=ows_client, impersonated_identity_uuid="some-uuid" ) class TestCalculateNumPartsAndPartSizeBytes: def test_calculate_part_size_under_max_parts(self) -> None: file_size_bytes = 10 * uploader._MIB result = uploader._calculate_num_parts_and_part_size_bytes(file_size_bytes) assert result == (1, uploader._PART_SIZE_BYTES) def test_calculate_part_size_at_max_parts(self) -> None: file_size_bytes = uploader._MAX_NUM_PARTS * uploader._PART_SIZE_BYTES result = uploader._calculate_num_parts_and_part_size_bytes(file_size_bytes) assert result == (uploader._MAX_NUM_PARTS, uploader._PART_SIZE_BYTES) def test_calculate_part_size_just_over_max_parts(self) -> None: file_size_bytes = (uploader._MAX_NUM_PARTS * uploader._PART_SIZE_BYTES) + 1 result = uploader._calculate_num_parts_and_part_size_bytes(file_size_bytes) assert result == (uploader._MAX_NUM_PARTS, uploader._PART_SIZE_BYTES + 1) def test_calculate_part_size_way_over_max_parts(self) -> None: file_size_bytes = 2 * (uploader._MAX_NUM_PARTS * uploader._PART_SIZE_BYTES) result = uploader._calculate_num_parts_and_part_size_bytes(file_size_bytes) assert result == (uploader._MAX_NUM_PARTS, 2 * uploader._PART_SIZE_BYTES) class TestGetSourceFilename: @pytest.mark.parametrize( ( "test_description", "filesystem_class", "path", "input_source_filename", "expected_result", ), [ ( "source filename provided", AbstractFileSystem, "/any/path/file.wav", "custom_name.wav", "custom_name.wav", ), ( "local file path", AbstractFileSystem, "/path/to/audio.wav", None, "audio.wav", ), ( "s3 path", S3FileSystem, "s3://bucket/folder/music.mp3", None, "music.mp3", ), ( "http path", HTTPFileSystem, "https://some-bucket.s3.us-east-1.amazonaws.com/0004d9c8_0086_4834_80cb_ab0ae9661d4b.wav?X-Amz-Signature=abc", None, "0004d9c8_0086_4834_80cb_ab0ae9661d4b.wav", ), ( "filename with no extension", AbstractFileSystem, "/path/to/audiofile", None, "audiofile", ), ( "http path with no extension", HTTPFileSystem, "https://some-bucket.s3.us-east-1.amazonaws.com/0004d9c8_0086_4834_80cb_ab0ae9661d4b?X-Amz-Signature=abc", None, "0004d9c8_0086_4834_80cb_ab0ae9661d4b", ), ], ) def test_get_source_filename( self, test_description: str, filesystem_class: type, path: str, input_source_filename: str | None, expected_result: str, ) -> None: filesystem = Mock(spec=filesystem_class) result = uploader._get_source_filename(filesystem, path, input_source_filename) assert result == expected_result class TestCreateMultipartUpload: @pytest.mark.parametrize( ("uploader_fixture_name", "should_check_impersonated_identity_uuid"), [ pytest.param("assets_uploader", False, id="standard client"), pytest.param("obo_assets_uploader", True, id="impersonation client"), ], ) @pytest.mark.parametrize( ( "source_filename", "asset_upload_type", "track_id", "expected_track_unique_id", ), [ pytest.param("audio.wav", "stereo", 456, 456, id="track asset"), pytest.param("image.tif", "static_artwork", None, 0, id="product asset"), ], ) def test_create_multipart_upload( self, request: pytest.FixtureRequest, ows_client_mock: OwsClientMock, uploader_fixture_name: str, should_check_impersonated_identity_uuid: bool, impersonated_identity_uuid: str, source_filename: str, asset_upload_type: str, track_id: int | None, expected_track_unique_id: int, ) -> None: assets_uploader: AssetsUploader = request.getfixturevalue(uploader_fixture_name) ows_client_mock.post( "ows-assets", "/v2/assets/upload", json={ "product_id": 123, "track_unique_id": expected_track_unique_id, "original_filename": source_filename, "asset_upload_type": asset_upload_type, }, headers={"Orchard-User-Id": "oa:179", "Content-Type": "application/json"}, ).mock( return_value=httpx.Response(200, json={"filename": "destination_filename"}) ) result = assets_uploader._create_multipart_upload( source_filename, asset_upload_type, 123, track_id ) assert result == "destination_filename" if should_check_impersonated_identity_uuid: ows_client_mock.mock_impersonation_m2m_token_manager.get_token_string.assert_called_once_with( # type: ignore[attr-defined] impersonated_identity_uuid=impersonated_identity_uuid, ) else: ows_client_mock.mock_impersonation_m2m_token_manager.get_token_string.assert_not_called() # type: ignore[attr-defined] class TestGetPartUploadUrlBatch: @pytest.mark.parametrize( ("uploader_fixture_name", "should_check_impersonated_identity_uuid"), [ pytest.param("assets_uploader", False, id="standard client"), pytest.param("obo_assets_uploader", True, id="impersonation client"), ], ) def test_get_part_upload_url_batch( self, request: pytest.FixtureRequest, ows_client_mock: OwsClientMock, uploader_fixture_name: str, should_check_impersonated_identity_uuid: bool, impersonated_identity_uuid: str, ) -> None: assets_uploader: AssetsUploader = request.getfixturevalue(uploader_fixture_name) ows_client_mock.get( "ows-assets", "/v2/assets/upload/test.txt", params={"part_numbers": "1,2"}, headers={"Orchard-User-Id": "oa:179"}, ).mock( return_value=httpx.Response( 200, json={ "part_number_to_presigned_url": { "1": "https://example.com/upload1", "2": "https://example.com/upload2", } }, ) ) result = assets_uploader._get_part_upload_url_batch("test.txt", [1, 2]) assert result == { 1: "https://example.com/upload1", 2: "https://example.com/upload2", } if should_check_impersonated_identity_uuid: ows_client_mock.mock_impersonation_m2m_token_manager.get_token_string.assert_called_once_with( # type: ignore[attr-defined] impersonated_identity_uuid=impersonated_identity_uuid, ) else: ows_client_mock.mock_impersonation_m2m_token_manager.get_token_string.assert_not_called() # type: ignore[attr-defined] class TestGeneratePartUploadUrls: @patch("ows_assets_uploader.uploader._MAX_PRESIGNED_URL_BATCH_SIZE", 3) def test_generate_part_upload_urls( self, assets_uploader: AssetsUploader, ) -> None: with patch.object( assets_uploader, "_get_part_upload_url_batch" ) as mock_get_batch: mock_get_batch.side_effect = [ { 1: "https://example.com/upload1", 2: "https://example.com/upload2", 3: "https://example.com/upload3", }, { 4: "https://example.com/upload4", 5: "https://example.com/upload5", }, ] result = list(assets_uploader._generate_part_upload_urls("test.txt", 5)) assert result == [ (1, "https://example.com/upload1"), (2, "https://example.com/upload2"), (3, "https://example.com/upload3"), (4, "https://example.com/upload4"), (5, "https://example.com/upload5"), ] class TestCompleteMultipartUpload: @pytest.mark.parametrize( ("uploader_fixture_name", "should_check_impersonated_identity_uuid"), [ pytest.param("assets_uploader", False, id="standard client"), pytest.param("obo_assets_uploader", True, id="impersonation client"), ], ) def test_complete_multipart_upload( self, request: pytest.FixtureRequest, ows_client_mock: OwsClientMock, uploader_fixture_name: str, should_check_impersonated_identity_uuid: bool, impersonated_identity_uuid: str, ) -> None: assets_uploader: AssetsUploader = request.getfixturevalue(uploader_fixture_name) completed_parts = [ uploader._CompletedPart(part_number=2, etag="etag2"), uploader._CompletedPart(part_number=1, etag="etag1"), uploader._CompletedPart(part_number=3, etag="etag3"), ] ows_client_mock.patch( "ows-assets", "/v2/assets/upload/destination_filename", json={ "parts": [ {"part_number": 1, "etag": "etag1"}, {"part_number": 2, "etag": "etag2"}, {"part_number": 3, "etag": "etag3"}, ] }, headers={"Orchard-User-Id": "oa:179", "Content-Type": "application/json"}, ).mock(return_value=httpx.Response(200)) assets_uploader._complete_multipart_upload( "destination_filename", completed_parts ) if should_check_impersonated_identity_uuid: ows_client_mock.mock_impersonation_m2m_token_manager.get_token_string.assert_called_once_with( # type: ignore[attr-defined] impersonated_identity_uuid=impersonated_identity_uuid, ) else: ows_client_mock.mock_impersonation_m2m_token_manager.get_token_string.assert_not_called() # type: ignore[attr-defined] class TestReadPartData: def test_read_part_data_first_part(self) -> None: file_bytes = b"0123456789ABCDEF" file = io.BytesIO(file_bytes) part_size = 5 file_size_bytes = len(file_bytes) num_parts = math.ceil(file_size_bytes / part_size) assert [ uploader._read_part_data(file, part_number, part_size, file_size_bytes) for part_number in range(1, num_parts + 1) ] == [b"01234", b"56789", b"ABCDE", b"F"] class TestUpload: @patch("ows_assets_uploader.uploader._PART_SIZE_BYTES", 5) def test_upload( self, assets_uploader: AssetsUploader, ) -> None: product_id = 456 track_id = 789 file_bytes = b"0123456789ABCDEF" file_size_bytes = len(file_bytes) part_size_bytes = 5 num_parts = math.ceil(file_size_bytes / part_size_bytes) with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as file: file.write(file_bytes) source_file_location = file.name source_filename = os.path.basename(source_file_location) destination_filename = "destination_filename" with ( patch.object( assets_uploader, "_create_multipart_upload" ) as mock_create_multipart_upload, patch.object( assets_uploader, "_generate_part_upload_urls" ) as mock_generate_part_upload_urls, patch.object( assets_uploader, "_complete_multipart_upload" ) as mock_complete_multipart_upload, ): mock_create_multipart_upload.return_value = destination_filename mock_generate_part_upload_urls.return_value = [ (1, "https://example.com/upload1"), (2, "https://example.com/upload2"), (3, "https://example.com/upload3"), (4, "https://example.com/upload4"), ] mock_httpx_client = Mock(spec=httpx.Client) mock_httpx_client.put.side_effect = [ Mock(headers={"ETag": "etag1"}), Mock(headers={"ETag": "etag2"}), Mock(headers={"ETag": "etag3"}), Mock(headers={"ETag": "etag4"}), ] assets_uploader._httpx_client = mock_httpx_client result = assets_uploader.upload( source_file_location=source_file_location, asset_upload_type="stereo", product_id=product_id, track_id=track_id, ) mock_create_multipart_upload.assert_called_once_with( source_filename, "stereo", product_id, track_id, ) mock_generate_part_upload_urls.assert_called_once_with( destination_filename, num_parts, ) assert mock_httpx_client.put.mock_calls == [ call("https://example.com/upload1", content=b"01234"), call("https://example.com/upload2", content=b"56789"), call("https://example.com/upload3", content=b"ABCDE"), call("https://example.com/upload4", content=b"F"), ] for mock_response in mock_httpx_client.put.side_effect: mock_response.raise_for_status.assert_called_once() mock_complete_multipart_upload.assert_called_once_with( destination_filename, [ uploader._CompletedPart(part_number=1, etag="etag1"), uploader._CompletedPart(part_number=2, etag="etag2"), uploader._CompletedPart(part_number=3, etag="etag3"), uploader._CompletedPart(part_number=4, etag="etag4"), ], ) assert result == destination_filename os.remove(source_file_location) @pytest.mark.parametrize( ("file_extension", "upload_kwargs", "expected_track_id"), [ pytest.param( ".tif", {"asset_upload_type": "static_artwork"}, None, id="product asset", ), pytest.param( ".wav", {"asset_upload_type": "stereo", "track_id": 456}, 456, id="track asset", ), ], ) def test_upload_forwards_args_to_create_multipart_upload( self, file_extension: str, upload_kwargs: dict[str, Any], expected_track_id: int | None, assets_uploader: AssetsUploader, ) -> None: with patch.object( assets_uploader, "_create_multipart_upload" ) as mock_create_multipart_upload: mock_create_multipart_upload.side_effect = Exception( "Stopping at _create_multipart_upload" ) product_id = 123 with tempfile.NamedTemporaryFile( suffix=file_extension, delete=False ) as file: source_file_location = file.name source_filename = os.path.basename(source_file_location) try: assets_uploader.upload( source_file_location=source_file_location, **upload_kwargs, product_id=product_id, ) except Exception: pass mock_create_multipart_upload.assert_called_once_with( source_filename, upload_kwargs["asset_upload_type"], product_id, expected_track_id, ) os.remove(source_file_location) class TestModuleLevelUpload: @patch("ows_assets_uploader.uploader.AssetsUploader") def test_upload_uses_provided_ows_client( self, mock_assets_uploader_cls: Mock ) -> None: mock_ows_client = Mock() uploader.upload( source_file_location="/tmp/test.wav", asset_upload_type="stereo", product_id=123, ows_client=mock_ows_client, ) mock_assets_uploader_cls.assert_called_once_with(ows_client=mock_ows_client) mock_assets_uploader_cls.return_value.upload.assert_called_once_with( source_file_location="/tmp/test.wav", asset_upload_type="stereo", product_id=123, track_id=None, source_filename=None, ) @patch("ows_assets_uploader.uploader.AssetsUploader") @patch("ows_assets_uploader.uploader._get_ows_client") def test_upload_uses_default_ows_client_when_none( self, mock_get_ows_client: Mock, mock_assets_uploader_cls: Mock ) -> None: uploader.upload( source_file_location="/tmp/test.wav", asset_upload_type="stereo", product_id=123, ) mock_get_ows_client.assert_called_once() mock_assets_uploader_cls.assert_called_once_with( ows_client=mock_get_ows_client.return_value )