"""Tests for OWS Assets API client.""" from unittest.mock import MagicMock from unittest.mock import patch import httpx import pytest from src.connectors import ows_assets @patch("src.connectors.ows_assets.client") def test_get_presigned_url_success(mock_ows_client): """Test get_presigned_url with a successful response.""" asset_final_id = 12345 mock_response = MagicMock() mock_response.status_code = 200 mock_response.json.return_value = {"url": "http://example.com/asset"} mock_ows_client.get.return_value = mock_response url = ows_assets.get_presigned_url(asset_final_id) assert url == "http://example.com/asset" mock_ows_client.get.assert_called_once_with( "ows-assets", "/v2/asset/download_url", params={ "asset_final_id": asset_final_id, "expires_in": 300, }, ) @patch("src.connectors.ows_assets.client") def test_get_presigned_url_http_error(mock_ows_client): """Test get_presigned_url with an HTTP error.""" asset_final_id = 12345 mock_response = MagicMock() mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( "Not Found", request=MagicMock(), response=MagicMock() ) mock_ows_client.get.return_value = mock_response with pytest.raises(httpx.HTTPStatusError): ows_assets.get_presigned_url(asset_final_id) mock_ows_client.get.assert_called_once_with( "ows-assets", "/v2/asset/download_url", params={ "asset_final_id": asset_final_id, "expires_in": 300, }, ) @pytest.fixture def hive_api_response_data(): return { "id": "52f803d0-aeb3-11f0-b444-87bc4223b971", "code": 200, "project_id": 1104290940, "user_id": 6661, "created_on": "2025-10-21T19:22:51.987Z", "status": [ { "status": {"code": "0", "message": "SUCCESS"}, "response": { "output": [ { "time": 0, "classes": [ { "class": "ai_generated_music", "score": 2.999325161043201e-07, }, { "class": "not_ai_generated_music", "score": 0.9999997000674838, }, { "class": "ai_generated_music_vocal", "score": 2.6973483452557593e-07, }, { "class": "not_ai_generated_music_vocal", "score": 0.9999997302651654, }, {"class": "riffusion", "score": 1.0112554402526732e-05}, { "class": "stable_audio", "score": 5.258721193160129e-06, }, {"class": "suno", "score": 2.272364926179113e-05}, {"class": "udio", "score": 3.3340098796857813e-06}, {"class": "mubert", "score": 1.681489804075235e-05}, {"class": "musicgen", "score": 2.9261321591240557e-06}, {"class": "yue", "score": 3.638046655156766e-05}, {"class": "minimax", "score": 3.725670280721739e-05}, {"class": "mureka", "score": 2.2318162129024265e-05}, {"class": "ace_step", "score": 1.111111111111111e-05}, {"class": "duobao", "score": 2.222222222222222e-05}, {"class": "google", "score": 3.333333333333333e-05}, {"class": "heartmula", "score": 4.444444444444444e-05}, {"class": "loudly", "score": 5.555555555555555e-05}, ], } ], "input": { "model_version": 1, "model": "ai_music_classifier_DORIAN_2025_04_02_v00", "inference_client_version": "0.0.0", "hash": "06a80cf032ecc3e0e8a5ffa3162773cf", "media": { "url": None, "filename": None, "type": "AUDIO", "mime_type": "x-flac", "mimetype": "audio/x-flac", "width": 0, "height": 0, "num_frames": 61, "duration": 183, }, "id": "52f803d0-aeb3-11f0-b444-87bc4223b971", "charge": 0.0006100000000000001, "created_on": "2025-10-21T19:22:46.285Z", "model_type": "AI_AUDIO_DETECTION", "user_id": 6661, "project_id": 1104290940, }, }, } ], "from_cache": False, "hsl_options": {}, } def test_format_raw_hive_segment_data(hive_api_response_data): """Test _format_raw_hive_segment_data function.""" formatted_data = ows_assets._format_raw_hive_segment_data(hive_api_response_data) assert formatted_data == [ { "time": 0, "ai_generated_music": 2.999325161043201e-07, "ai_generated_music_vocal": 2.6973483452557593e-07, "mubert": 1.681489804075235e-05, "musicgen": 2.9261321591240557e-06, "riffusion": 1.0112554402526732e-05, "suno": 2.272364926179113e-05, "stable_audio": 5.258721193160129e-06, "udio": 3.3340098796857813e-06, "yue": 3.638046655156766e-05, "minimax": 3.725670280721739e-05, "mureka": 2.2318162129024265e-05, "ace_step": 1.111111111111111e-05, "duobao": 2.222222222222222e-05, "google": 3.333333333333333e-05, "heartmula": 4.444444444444444e-05, "loudly": 5.555555555555555e-05, } ] def test_format_raw_hive_segment_data_remaps_stable_audio_open(hive_api_response_data): """stable_audio_open from HYMN models is stored under the stable_audio column.""" classes = hive_api_response_data["status"][0]["response"]["output"][0]["classes"] for entry in classes: if entry["class"] == "stable_audio": entry["class"] = "stable_audio_open" formatted_data = ows_assets._format_raw_hive_segment_data(hive_api_response_data) assert formatted_data[0]["stable_audio"] == 5.258721193160129e-06 assert "stable_audio_open" not in formatted_data[0] def test_format_raw_hive_task_data(hive_api_response_data): """Test _format_raw_hive_task_data function.""" formatted_task_data = ows_assets._format_raw_hive_task_data(hive_api_response_data) assert formatted_task_data == { "task_id": "52f803d0-aeb3-11f0-b444-87bc4223b971", "model": "ai_music_classifier_DORIAN_2025_04_02_v00", "model_version": 1, } @patch("src.connectors.ows_assets.client") @patch("src.connectors.ows_assets._format_raw_hive_segment_data") @patch("src.connectors.ows_assets._format_raw_hive_task_data") def test_post_results_success(mock_task_format, mock_segment_format, mock_ows_client): """Test post_results with a successful response.""" hive_segment_data = {"data": "some_data"} hive_task_data = {"data": "test_task_data"} mock_segment_format.return_value = hive_segment_data mock_task_format.return_value = hive_task_data mock_response = MagicMock() mock_response.status_code = 200 mock_ows_client.post.return_value = mock_response ows_assets.post_results(12345, {}) mock_ows_client.post.assert_called_once_with( "ows-assets", "/hive_segments", json={ "asset_final_id": 12345, "task": hive_task_data, "segments": hive_segment_data, }, ) @patch("src.connectors.ows_assets.client") @patch("src.connectors.ows_assets._format_raw_hive_segment_data") @patch("src.connectors.ows_assets._format_raw_hive_task_data") def test_post_results_overwrite(mock_task_format, mock_segment_format, mock_ows_client): """overwrite=True adds the flag so the endpoint replaces any existing scan.""" hive_segment_data = {"data": "some_data"} hive_task_data = {"data": "test_task_data"} mock_segment_format.return_value = hive_segment_data mock_task_format.return_value = hive_task_data mock_ows_client.post.return_value = MagicMock(status_code=200) ows_assets.post_results(12345, {}, overwrite=True) mock_ows_client.post.assert_called_once_with( "ows-assets", "/hive_segments", json={ "asset_final_id": 12345, "task": hive_task_data, "segments": hive_segment_data, "overwrite": True, }, )