test: add unit tests for WebSocketOptions and its integration in streaming

Author: twangodev

tests/integration/test_tts_websocket_integration.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from fishaudio import WebSocketOptions
 from fishaudio.types import Prosody, TTSConfig, TextEvent, FlushEvent
 from .conftest import TEST_REFERENCE_ID
 
@@ -118,6 +119,59 @@ def text_stream():
         with pytest.raises(WebSocketError, match="WebSocket stream ended with error"):
             list(client.tts.stream_websocket(text_stream()))
 
+    def test_websocket_very_long_generation_with_timeout(self, client, save_audio):
+        """
+        Test that very long text generation succeeds with increased timeout.
+
+        This test generates a very long response that could potentially take >20 seconds
+        to fully generate, which would cause a WebSocketNetworkError with the default
+        keepalive_ping_timeout_seconds=20. By using an increased timeout of 60 seconds,
+        we can handle longer generation times without disconnection.
+
+        This is the SOLUTION to issue #47. To reproduce the timeout issue, run:
+        python reproduce_issue_47.py --mode=both
+        """
+        # Use significantly increased timeout to handle very long generations
+        ws_options = WebSocketOptions(
+            keepalive_ping_timeout_seconds=60.0,
+            keepalive_ping_interval_seconds=30.0,
+        )
+
+        def text_stream():
+            # Generate a very long piece of text that will take significant time to process
+            long_text = [
+                "This is a test of very long form text-to-speech generation. ",
+                "We are testing the ability to handle extended generation times without timing out. ",
+                "The default WebSocket keepalive timeout of 20 seconds can be insufficient for long responses. ",
+                "By increasing the keepalive_ping_timeout_seconds to 60 seconds, we allow for longer gaps between chunks. ",
+                "This is particularly important for conversational AI applications where responses can be quite lengthy. ",
+                "The WebSocket connection should remain stable throughout the entire generation process. ",
+                "We include enough text here to ensure the generation takes a substantial amount of time. ",
+                "This helps verify that the increased timeout setting is working correctly. ",
+                "The audio streaming should continue smoothly without any network errors. ",
+                "Each sentence adds more content to be synthesized into speech. ",
+                "The system should handle this gracefully with the custom WebSocket options. ",
+                "This demonstrates the practical value of the WebSocketOptions feature. ",
+                "Users can now configure timeouts based on their specific use case requirements. ",
+                "Long-form content generation is now much more reliable. ",
+                "The implementation passes through all necessary parameters to the underlying httpx_ws library. ",
+            ]
+            for sentence in long_text:
+                yield sentence
+
+        # This should succeed with increased timeout
+        audio_chunks = list(
+            client.tts.stream_websocket(text_stream(), ws_options=ws_options)
+        )
+
+        assert len(audio_chunks) > 0, "Should receive audio chunks for very long text"
+        complete_audio = b"".join(audio_chunks)
+        # Very long text should produce substantial audio
+        assert len(complete_audio) > 10000, (
+            "Very long text should produce substantial audio data"
+        )
+        save_audio(audio_chunks, "test_websocket_very_long_with_timeout.mp3")
+
 
 class TestAsyncTTSWebSocketIntegration:
     """Test async TTS WebSocket streaming with real API."""

tests/unit/test_core.py

@@ -4,7 +4,13 @@
 from unittest.mock import patch
 import httpx
 
-from fishaudio.core import OMIT, ClientWrapper, AsyncClientWrapper, RequestOptions
+from fishaudio.core import (
+    OMIT,
+    ClientWrapper,
+    AsyncClientWrapper,
+    RequestOptions,
+    WebSocketOptions,
+)
 
 
 class TestOMIT:
@@ -51,6 +57,38 @@ def test_get_timeout(self):
         assert timeout.connect == 30.0
 
 
+class TestWebSocketOptions:
+    """Test WebSocketOptions class."""
+
+    def test_to_httpx_ws_kwargs_all_options(self):
+        """Test to_httpx_ws_kwargs with all options set."""
+        options = WebSocketOptions(
+            keepalive_ping_timeout_seconds=60.0,
+            keepalive_ping_interval_seconds=30.0,
+            max_message_size_bytes=131072,
+            queue_size=1024,
+        )
+        kwargs = options.to_httpx_ws_kwargs()
+        assert kwargs == {
+            "keepalive_ping_timeout_seconds": 60.0,
+            "keepalive_ping_interval_seconds": 30.0,
+            "max_message_size_bytes": 131072,
+            "queue_size": 1024,
+        }
+
+    def test_to_httpx_ws_kwargs_partial_options(self):
+        """Test to_httpx_ws_kwargs with only some options set."""
+        options = WebSocketOptions(keepalive_ping_timeout_seconds=60.0)
+        kwargs = options.to_httpx_ws_kwargs()
+        assert kwargs == {"keepalive_ping_timeout_seconds": 60.0}
+        assert "keepalive_ping_interval_seconds" not in kwargs
+
+    def test_to_httpx_ws_kwargs_no_options(self):
+        """Test to_httpx_ws_kwargs with no options set."""
+        options = WebSocketOptions()
+        assert options.to_httpx_ws_kwargs() == {}
+
+
 class TestClientWrapper:
     """Test sync ClientWrapper."""
 

tests/unit/test_tts_realtime.py

@@ -3,7 +3,7 @@
 import pytest
 from unittest.mock import Mock, AsyncMock, MagicMock, patch
 
-from fishaudio.core import ClientWrapper, AsyncClientWrapper
+from fishaudio.core import ClientWrapper, AsyncClientWrapper, WebSocketOptions
 from fishaudio.resources.tts import TTSClient, AsyncTTSClient
 from fishaudio.types import Prosody, TTSConfig, TextEvent, FlushEvent, ReferenceAudio
 import ormsgpack
@@ -345,6 +345,30 @@ def submit_side_effect(fn):
             assert len(start_event_payload["request"]["references"]) == 1
             assert start_event_payload["request"]["references"][0]["text"] == "Param"
 
+    @patch("fishaudio.resources.tts.connect_ws")
+    @patch("fishaudio.resources.tts.ThreadPoolExecutor")
+    def test_stream_websocket_with_ws_options(
+        self, mock_executor, mock_connect_ws, tts_client, mock_client_wrapper
+    ):
+        """Test WebSocket streaming passes through WebSocketOptions."""
+        mock_ws = MagicMock()
+        mock_ws.__enter__ = Mock(return_value=mock_ws)
+        mock_ws.__exit__ = Mock(return_value=None)
+        mock_connect_ws.return_value = mock_ws
+        mock_future = Mock()
+        mock_future.result.return_value = None
+        mock_executor_instance = Mock()
+        mock_executor_instance.submit.return_value = mock_future
+        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"])
+            ws_options = WebSocketOptions(keepalive_ping_timeout_seconds=60.0)
+            list(tts_client.stream_websocket(iter(["Test"]), ws_options=ws_options))
+            assert (
+                mock_connect_ws.call_args[1]["keepalive_ping_timeout_seconds"] == 60.0
+            )
+
 
 class TestAsyncTTSRealtimeClient:
     """Test asynchronous AsyncTTSClient realtime streaming."""
@@ -649,3 +673,35 @@ async def text_stream():
             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"] == "Param"
+
+    @pytest.mark.asyncio
+    @patch("fishaudio.resources.tts.aconnect_ws")
+    async def test_stream_websocket_with_ws_options(
+        self, mock_aconnect_ws, async_tts_client, async_mock_client_wrapper
+    ):
+        """Test async WebSocket streaming passes through WebSocketOptions."""
+        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),
+        ):
+            ws_options = WebSocketOptions(keepalive_ping_timeout_seconds=60.0)
+
+            async def text_stream():
+                yield "Test"
+
+            async for _ in async_tts_client.stream_websocket(
+                text_stream(), ws_options=ws_options
+            ):
+                pass
+            assert (
+                mock_aconnect_ws.call_args[1]["keepalive_ping_timeout_seconds"] == 60.0
+            )