feat: add tests for WebSocket streaming with reference_id and references parameters

Signed-off-by: James Ding <[email protected]>

Author: twangodev

tests/unit/test_tts_realtime.py

@@ -5,7 +5,8 @@
 
 from fishaudio.core import ClientWrapper, AsyncClientWrapper
 from fishaudio.resources.tts import TTSClient, AsyncTTSClient
-from fishaudio.types import Prosody, TTSConfig, TextEvent, FlushEvent
+from fishaudio.types import Prosody, TTSConfig, TextEvent, FlushEvent, ReferenceAudio
+import ormsgpack
 
 
 @pytest.fixture
@@ -181,6 +182,169 @@ def test_stream_websocket_max_workers(
             # Verify ThreadPoolExecutor was created with max_workers=5
             mock_executor.assert_called_once_with(max_workers=5)
 
+    @patch("fishaudio.resources.tts.connect_ws")
+    @patch("fishaudio.resources.tts.ThreadPoolExecutor")
+    def test_stream_websocket_with_reference_id_parameter(
+        self, mock_executor, mock_connect_ws, tts_client, mock_client_wrapper
+    ):
+        """Test WebSocket streaming with reference_id as direct parameter."""
+        # Setup mocks
+        mock_ws = MagicMock()
+        mock_ws.__enter__ = Mock(return_value=mock_ws)
+        mock_ws.__exit__ = Mock(return_value=None)
+        mock_ws.send_bytes = Mock()
+        mock_connect_ws.return_value = mock_ws
+
+        # Make executor.submit actually run the function
+        def submit_side_effect(fn):
+            fn()  # Execute the sender function
+            mock_future = Mock()
+            mock_future.result.return_value = None
+            return mock_future
+
+        mock_executor_instance = Mock()
+        mock_executor_instance.submit.side_effect = submit_side_effect
+        mock_executor.return_value = mock_executor_instance
+
+        with patch("fishaudio.resources.tts.iter_websocket_audio") as mock_receiver:
+            mock_receiver.return_value = iter([b"audio"])
+
+            text_stream = iter(["Test"])
+            list(tts_client.stream_websocket(text_stream, reference_id="voice_456"))
+
+            # Verify WebSocket was called with StartEvent containing reference_id
+            assert mock_ws.send_bytes.called
+            # Get the first call (StartEvent)
+            first_call = mock_ws.send_bytes.call_args_list[0]
+            start_event_payload = ormsgpack.unpackb(first_call[0][0])
+            assert start_event_payload["request"]["reference_id"] == "voice_456"
+
+    @patch("fishaudio.resources.tts.connect_ws")
+    @patch("fishaudio.resources.tts.ThreadPoolExecutor")
+    def test_stream_websocket_config_reference_id_overrides_parameter(
+        self, mock_executor, mock_connect_ws, tts_client, mock_client_wrapper
+    ):
+        """Test that config.reference_id overrides parameter reference_id."""
+        # Setup mocks
+        mock_ws = MagicMock()
+        mock_ws.__enter__ = Mock(return_value=mock_ws)
+        mock_ws.__exit__ = Mock(return_value=None)
+        mock_ws.send_bytes = Mock()
+        mock_connect_ws.return_value = mock_ws
+
+        # Make executor.submit actually run the function
+        def submit_side_effect(fn):
+            fn()  # Execute the sender function
+            mock_future = Mock()
+            mock_future.result.return_value = None
+            return mock_future
+
+        mock_executor_instance = Mock()
+        mock_executor_instance.submit.side_effect = submit_side_effect
+        mock_executor.return_value = mock_executor_instance
+
+        with patch("fishaudio.resources.tts.iter_websocket_audio") as mock_receiver:
+            mock_receiver.return_value = iter([b"audio"])
+
+            config = TTSConfig(reference_id="voice_from_config")
+            text_stream = iter(["Test"])
+            list(
+                tts_client.stream_websocket(
+                    text_stream, reference_id="voice_from_param", config=config
+                )
+            )
+
+            # Verify config reference_id takes precedence
+            first_call = mock_ws.send_bytes.call_args_list[0]
+            start_event_payload = ormsgpack.unpackb(first_call[0][0])
+            assert start_event_payload["request"]["reference_id"] == "voice_from_config"
+
+    @patch("fishaudio.resources.tts.connect_ws")
+    @patch("fishaudio.resources.tts.ThreadPoolExecutor")
+    def test_stream_websocket_with_references_parameter(
+        self, mock_executor, mock_connect_ws, tts_client, mock_client_wrapper
+    ):
+        """Test WebSocket streaming with references as direct parameter."""
+        # Setup mocks
+        mock_ws = MagicMock()
+        mock_ws.__enter__ = Mock(return_value=mock_ws)
+        mock_ws.__exit__ = Mock(return_value=None)
+        mock_ws.send_bytes = Mock()
+        mock_connect_ws.return_value = mock_ws
+
+        # Make executor.submit actually run the function
+        def submit_side_effect(fn):
+            fn()  # Execute the sender function
+            mock_future = Mock()
+            mock_future.result.return_value = None
+            return mock_future
+
+        mock_executor_instance = Mock()
+        mock_executor_instance.submit.side_effect = submit_side_effect
+        mock_executor.return_value = mock_executor_instance
+
+        with patch("fishaudio.resources.tts.iter_websocket_audio") as mock_receiver:
+            mock_receiver.return_value = iter([b"audio"])
+
+            references = [
+                ReferenceAudio(audio=b"ref_audio_1", text="Sample 1"),
+                ReferenceAudio(audio=b"ref_audio_2", text="Sample 2"),
+            ]
+
+            text_stream = iter(["Test"])
+            list(tts_client.stream_websocket(text_stream, references=references))
+
+            # Verify references in StartEvent
+            first_call = mock_ws.send_bytes.call_args_list[0]
+            start_event_payload = ormsgpack.unpackb(first_call[0][0])
+            assert len(start_event_payload["request"]["references"]) == 2
+            assert start_event_payload["request"]["references"][0]["text"] == "Sample 1"
+            assert start_event_payload["request"]["references"][1]["text"] == "Sample 2"
+
+    @patch("fishaudio.resources.tts.connect_ws")
+    @patch("fishaudio.resources.tts.ThreadPoolExecutor")
+    def test_stream_websocket_config_references_overrides_parameter(
+        self, mock_executor, mock_connect_ws, tts_client, mock_client_wrapper
+    ):
+        """Test that config.references overrides parameter references."""
+        # Setup mocks
+        mock_ws = MagicMock()
+        mock_ws.__enter__ = Mock(return_value=mock_ws)
+        mock_ws.__exit__ = Mock(return_value=None)
+        mock_ws.send_bytes = Mock()
+        mock_connect_ws.return_value = mock_ws
+
+        # Make executor.submit actually run the function
+        def submit_side_effect(fn):
+            fn()  # Execute the sender function
+            mock_future = Mock()
+            mock_future.result.return_value = None
+            return mock_future
+
+        mock_executor_instance = Mock()
+        mock_executor_instance.submit.side_effect = submit_side_effect
+        mock_executor.return_value = mock_executor_instance
+
+        with patch("fishaudio.resources.tts.iter_websocket_audio") as mock_receiver:
+            mock_receiver.return_value = iter([b"audio"])
+
+            config_refs = [ReferenceAudio(audio=b"config_audio", text="Config")]
+            param_refs = [ReferenceAudio(audio=b"param_audio", text="Param")]
+
+            config = TTSConfig(references=config_refs)
+            text_stream = iter(["Test"])
+            list(
+                tts_client.stream_websocket(
+                    text_stream, references=param_refs, config=config
+                )
+            )
+
+            # Verify config references take precedence
+            first_call = mock_ws.send_bytes.call_args_list[0]
+            start_event_payload = ormsgpack.unpackb(first_call[0][0])
+            assert len(start_event_payload["request"]["references"]) == 1
+            assert start_event_payload["request"]["references"][0]["text"] == "Config"
+
 
 class TestAsyncTTSRealtimeClient:
     """Test asynchronous AsyncTTSClient realtime streaming."""
@@ -331,3 +495,157 @@ async def text_stream():
 
             # Should have no audio
             assert audio_chunks == []
+
+    @pytest.mark.asyncio
+    @patch("fishaudio.resources.tts.aconnect_ws")
+    async def test_stream_websocket_with_reference_id_parameter(
+        self, mock_aconnect_ws, async_tts_client, async_mock_client_wrapper
+    ):
+        """Test async WebSocket streaming with reference_id as direct parameter."""
+        # Setup mocks
+        mock_ws = MagicMock()
+        mock_ws.__aenter__ = AsyncMock(return_value=mock_ws)
+        mock_ws.__aexit__ = AsyncMock(return_value=None)
+        mock_ws.send_bytes = AsyncMock()
+        mock_aconnect_ws.return_value = mock_ws
+
+        async def mock_audio_receiver(ws):
+            yield b"audio"
+
+        with patch(
+            "fishaudio.resources.tts.aiter_websocket_audio",
+            return_value=mock_audio_receiver(mock_ws),
+        ):
+
+            async def text_stream():
+                yield "Test"
+
+            audio_chunks = []
+            async for chunk in async_tts_client.stream_websocket(
+                text_stream(), reference_id="voice_456"
+            ):
+                audio_chunks.append(chunk)
+
+            # Verify WebSocket was called with StartEvent containing reference_id
+            assert mock_ws.send_bytes.called
+            # Get the first call (StartEvent)
+            first_call = mock_ws.send_bytes.call_args_list[0]
+            start_event_payload = ormsgpack.unpackb(first_call[0][0])
+            assert start_event_payload["request"]["reference_id"] == "voice_456"
+
+    @pytest.mark.asyncio
+    @patch("fishaudio.resources.tts.aconnect_ws")
+    async def test_stream_websocket_config_reference_id_overrides_parameter(
+        self, mock_aconnect_ws, async_tts_client, async_mock_client_wrapper
+    ):
+        """Test that config.reference_id overrides parameter reference_id (async)."""
+        # Setup mocks
+        mock_ws = MagicMock()
+        mock_ws.__aenter__ = AsyncMock(return_value=mock_ws)
+        mock_ws.__aexit__ = AsyncMock(return_value=None)
+        mock_ws.send_bytes = AsyncMock()
+        mock_aconnect_ws.return_value = mock_ws
+
+        async def mock_audio_receiver(ws):
+            yield b"audio"
+
+        with patch(
+            "fishaudio.resources.tts.aiter_websocket_audio",
+            return_value=mock_audio_receiver(mock_ws),
+        ):
+            config = TTSConfig(reference_id="voice_from_config")
+
+            async def text_stream():
+                yield "Test"
+
+            audio_chunks = []
+            async for chunk in async_tts_client.stream_websocket(
+                text_stream(), reference_id="voice_from_param", config=config
+            ):
+                audio_chunks.append(chunk)
+
+            # Verify config reference_id takes precedence
+            first_call = mock_ws.send_bytes.call_args_list[0]
+            start_event_payload = ormsgpack.unpackb(first_call[0][0])
+            assert start_event_payload["request"]["reference_id"] == "voice_from_config"
+
+    @pytest.mark.asyncio
+    @patch("fishaudio.resources.tts.aconnect_ws")
+    async def test_stream_websocket_with_references_parameter(
+        self, mock_aconnect_ws, async_tts_client, async_mock_client_wrapper
+    ):
+        """Test async WebSocket streaming with references as direct parameter."""
+        # Setup mocks
+        mock_ws = MagicMock()
+        mock_ws.__aenter__ = AsyncMock(return_value=mock_ws)
+        mock_ws.__aexit__ = AsyncMock(return_value=None)
+        mock_ws.send_bytes = AsyncMock()
+        mock_aconnect_ws.return_value = mock_ws
+
+        async def mock_audio_receiver(ws):
+            yield b"audio"
+
+        with patch(
+            "fishaudio.resources.tts.aiter_websocket_audio",
+            return_value=mock_audio_receiver(mock_ws),
+        ):
+            references = [
+                ReferenceAudio(audio=b"ref_audio_1", text="Sample 1"),
+                ReferenceAudio(audio=b"ref_audio_2", text="Sample 2"),
+            ]
+
+            async def text_stream():
+                yield "Test"
+
+            audio_chunks = []
+            async for chunk in async_tts_client.stream_websocket(
+                text_stream(), references=references
+            ):
+                audio_chunks.append(chunk)
+
+            # Verify references in StartEvent
+            first_call = mock_ws.send_bytes.call_args_list[0]
+            start_event_payload = ormsgpack.unpackb(first_call[0][0])
+            assert len(start_event_payload["request"]["references"]) == 2
+            assert start_event_payload["request"]["references"][0]["text"] == "Sample 1"
+            assert start_event_payload["request"]["references"][1]["text"] == "Sample 2"
+
+    @pytest.mark.asyncio
+    @patch("fishaudio.resources.tts.aconnect_ws")
+    async def test_stream_websocket_config_references_overrides_parameter(
+        self, mock_aconnect_ws, async_tts_client, async_mock_client_wrapper
+    ):
+        """Test that config.references overrides parameter references (async)."""
+        # Setup mocks
+        mock_ws = MagicMock()
+        mock_ws.__aenter__ = AsyncMock(return_value=mock_ws)
+        mock_ws.__aexit__ = AsyncMock(return_value=None)
+        mock_ws.send_bytes = AsyncMock()
+        mock_aconnect_ws.return_value = mock_ws
+
+        async def mock_audio_receiver(ws):
+            yield b"audio"
+
+        with patch(
+            "fishaudio.resources.tts.aiter_websocket_audio",
+            return_value=mock_audio_receiver(mock_ws),
+        ):
+            config_refs = [ReferenceAudio(audio=b"config_audio", text="Config")]
+            param_refs = [ReferenceAudio(audio=b"param_audio", text="Param")]
+
+            config = TTSConfig(references=config_refs)
+
+            async def text_stream():
+                yield "Test"
+
+            audio_chunks = []
+            async for chunk in async_tts_client.stream_websocket(
+                text_stream(), references=param_refs, config=config
+            ):
+                audio_chunks.append(chunk)
+
+            # Verify config references take precedence
+            first_call = mock_ws.send_bytes.call_args_list[0]
+            start_event_payload = ormsgpack.unpackb(first_call[0][0])
+            assert len(start_event_payload["request"]["references"]) == 1
+            assert start_event_payload["request"]["references"][0]["text"] == "Config"