diff --git a/tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py b/tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py index 8871260813d..16d8a033321 100644 --- a/tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py +++ b/tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py @@ -65,3 +65,88 @@ def test_runwayml_native_voice_passthrough(): assert "runwayml_voice" in mapped_params assert mapped_params["runwayml_voice"]["type"] == "runway-preset" assert mapped_params["runwayml_voice"]["presetId"] == runway_voice + +def test_transform_text_to_speech_response(): + """Test TTS audio download with SSRF-protected fetch.""" + from unittest.mock import Mock, patch + + import httpx + + from litellm.types.llms.openai import HttpxBinaryResponseContent + + config = RunwayMLTextToSpeechConfig() + + # Mock the initial response (task created) + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = {"id": "task-123", "status": "PENDING"} + mock_response.request.headers = {"Authorization": "Bearer test"} + + # Mock the polled response (task completed) + mock_polled = Mock(spec=httpx.Response) + mock_polled.json.return_value = { + "id": "task-123", + "status": "SUCCEEDED", + "output": ["https://example.com/audio.mp3"], + } + + # Mock the audio download response + mock_audio_response = Mock(spec=httpx.Response) + mock_audio_response.raise_for_status = Mock() + + with patch.object(config, "_poll_task_sync", return_value=mock_polled): + with patch("litellm.llms.custom_httpx.http_handler._get_httpx_client"): + with patch("litellm.llms.runwayml.text_to_speech.transformation.safe_get") as mock_safe_get: + mock_safe_get.return_value = mock_audio_response + result = config.transform_text_to_speech_response( + model="eleven_multilingual_v2", + raw_response=mock_response, + logging_obj=Mock(), + ) + + assert isinstance(result, HttpxBinaryResponseContent) + mock_safe_get.assert_called_once() + + +def test_async_transform_text_to_speech_response(): + """Test async TTS audio download with SSRF-protected fetch.""" + import asyncio + from unittest.mock import AsyncMock, Mock, patch + + import httpx + + from litellm.types.llms.openai import HttpxBinaryResponseContent + + config = RunwayMLTextToSpeechConfig() + + # Mock the initial response (task created) + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = {"id": "task-123", "status": "PENDING"} + mock_response.request.headers = {"Authorization": "Bearer test"} + + # Mock the polled response (task completed) + mock_polled = Mock(spec=httpx.Response) + mock_polled.json.return_value = { + "id": "task-123", + "status": "SUCCEEDED", + "output": ["https://example.com/audio.mp3"], + } + + # Mock the audio download response + mock_audio_response = Mock(spec=httpx.Response) + mock_audio_response.raise_for_status = Mock() + + async def run_test(): + with patch.object(config, "_poll_task_async", new_callable=AsyncMock, return_value=mock_polled): + with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"): + with patch("litellm.llms.runwayml.text_to_speech.transformation.async_safe_get", new_callable=AsyncMock) as mock_safe_get: + mock_safe_get.return_value = mock_audio_response + result = await config.async_transform_text_to_speech_response( + model="eleven_multilingual_v2", + raw_response=mock_response, + logging_obj=Mock(), + ) + return result + + result = asyncio.run(run_test()) + assert isinstance(result, HttpxBinaryResponseContent) + diff --git a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py b/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py index 24879ce83f9..e235ab7448d 100644 --- a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py +++ b/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py @@ -213,5 +213,62 @@ class TestRunwayMLVideoTransformation: assert isinstance(status_obj.completed_at, int) + + def test_transform_video_content_response(self): + """Test video content download with SSRF-protected fetch.""" + from unittest.mock import patch + + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "id": "test-id", + "status": "SUCCEEDED", + "output": ["https://example.com/video.mp4"], + } + + mock_video_response = Mock(spec=httpx.Response) + mock_video_response.content = b"fake-video-bytes" + mock_video_response.raise_for_status = Mock() + + with patch("litellm.llms.runwayml.videos.transformation._get_httpx_client") as mock_client: + with patch("litellm.llms.runwayml.videos.transformation.safe_get") as mock_safe_get: + mock_safe_get.return_value = mock_video_response + result = self.config.transform_video_content_response( + raw_response=mock_response, + logging_obj=self.mock_logging_obj, + ) + + assert result == b"fake-video-bytes" + mock_safe_get.assert_called_once() + + def test_async_transform_video_content_response(self): + """Test async video content download with SSRF-protected fetch.""" + import asyncio + from unittest.mock import AsyncMock, patch + + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "id": "test-id", + "status": "SUCCEEDED", + "output": ["https://example.com/video.mp4"], + } + + mock_video_response = Mock(spec=httpx.Response) + mock_video_response.content = b"fake-video-bytes" + mock_video_response.raise_for_status = Mock() + + async def run_test(): + with patch("litellm.llms.runwayml.videos.transformation.get_async_httpx_client") as mock_client: + with patch("litellm.llms.runwayml.videos.transformation.async_safe_get", new_callable=AsyncMock) as mock_safe_get: + mock_safe_get.return_value = mock_video_response + result = await self.config.async_transform_video_content_response( + raw_response=mock_response, + logging_obj=self.mock_logging_obj, + ) + return result + + result = asyncio.run(run_test()) + assert result == b"fake-video-bytes" + + if __name__ == "__main__": pytest.main([__file__, "-v"])