fix: enhance RedisStreamInjector to support custom Redis client and improve error handling

Author: mjunaidca

src/a2a/server/events/redis_queue_manager.py

@@ -36,7 +36,7 @@ def __init__(
     async def add(self, task_id: str, queue: EventQueue) -> None:
         """Add is not supported in distributed Redis setup.
 
-        In a distributed environment, we can't reliably add pre-existing queue
+        In a distributed environment, we can't reliably add preexisting queue
         instances. Use create_or_tap() instead to create Redis-backed queues.
         """
         raise NotImplementedError(

src/a2a/utils/stream_write/redis_stream_writer.py

@@ -26,25 +26,35 @@
 class RedisStreamInjector:
     """Professional stream injector for A2A framework."""
 
-    def __init__(self, redis_url: str = 'redis://localhost:6379/0'):
+    def __init__(
+        self,
+        redis_url: str = 'redis://localhost:6379/0',
+        redis_client: Any | None = None,
+    ):
         """Initialize the stream injector."""
-        if Redis is None:
+        # Allow passing a custom redis client (e.g. a fake in tests).
+        if Redis is None and redis_client is None:
             raise ImportError(
                 'redis package is required. Install with: pip install redis'
             )
 
         self.redis_url = redis_url
-        self._client = None
-        self._connected = False
+        self._client = redis_client
+        self._connected = redis_client is not None
 
     async def connect(self) -> None:
         """Establish Redis connection."""
         if self._connected:
             return
 
         try:
-            self._client = Redis.from_url(self.redis_url)
-            await self._client.ping()
+            if self._client is None:
+                if Redis is None:
+                    raise ImportError(
+                        'redis package is required. Install with: pip install redis'
+                    )
+                self._client = Redis.from_url(self.redis_url)
+                await self._client.ping()
             self._connected = True
             logger.info('Connected to Redis')
         except Exception:
@@ -102,7 +112,7 @@ async def _append_to_stream(
             raise RuntimeError('Not connected to Redis. Call connect() first.')
 
         stream_key = self._get_stream_key(task_id)
-        return await self._client.xadd(stream_key, event_data)
+        return await self._client.xadd(stream_key, event_data)  # type: ignore
 
     async def stream_message(
         self, context_id: str, task_id: str, message: dict[str, Any] | Message

tests/utils/test_redis_stream_writer.py

@@ -1,11 +1,20 @@
-import json
 import pytest
 from unittest.mock import AsyncMock, patch
 
 from a2a.types import TaskStatusUpdateEvent
 from a2a.utils.stream_write.redis_stream_writer import RedisStreamInjector
 
 
+@pytest.fixture
+def mock_redis_client():
+    """Fixture providing a mock Redis client."""
+    client = AsyncMock()
+    client.xadd = AsyncMock(return_value='123-0')
+    client.ping = AsyncMock()
+    client.aclose = AsyncMock()
+    return client
+
+
 class TestRedisStreamInjector:
     """Test suite for RedisStreamInjector."""
 
@@ -90,8 +99,7 @@ async def test_disconnect(self):
         mock_client = AsyncMock()
         mock_client.aclose = AsyncMock()
 
-        injector = RedisStreamInjector()
-        injector._client = mock_client
+        injector = RedisStreamInjector(redis_client=mock_client)
         injector._connected = True
 
         await injector.disconnect()
@@ -103,15 +111,18 @@ async def test_disconnect(self):
     @pytest.mark.asyncio
     async def test_disconnect_not_connected(self):
         """Test disconnect when not connected."""
-        injector = RedisStreamInjector()
-        injector._client = None
+        mock_client = AsyncMock()
+        injector = RedisStreamInjector(redis_client=mock_client)
         injector._connected = False
 
         await injector.disconnect()
 
-        # Should not raise any errors
+        # Should not call aclose since not connected
+        mock_client.aclose.assert_not_called()
+
+        # Should not raise any errors and client should remain
         assert not injector._connected
-        assert injector._client is None
+        assert injector._client == mock_client
 
     @pytest.mark.asyncio
     async def test_context_manager(self):
@@ -120,57 +131,46 @@ async def test_context_manager(self):
         mock_client.ping = AsyncMock()
         mock_client.aclose = AsyncMock()
 
-        with patch(
-            'a2a.utils.stream_write.redis_stream_writer.Redis'
-        ) as mock_redis_class:
-            mock_redis_class.from_url.return_value = mock_client
-
-            injector = RedisStreamInjector()
+        injector = RedisStreamInjector(redis_client=mock_client)
 
-            async with injector as ctx_injector:
-                assert ctx_injector == injector
-                assert injector._connected
+        async with injector as ctx_injector:
+            assert ctx_injector == injector
+            assert injector._connected
 
-            assert not injector._connected
-            mock_client.aclose.assert_called_once()
+        assert not injector._connected
+        mock_client.aclose.assert_called_once()
 
     def test_get_stream_key(self):
         """Test stream key generation."""
-        injector = RedisStreamInjector()
+        mock_client = AsyncMock()
+        injector = RedisStreamInjector(redis_client=mock_client)
 
         key = injector._get_stream_key('test_task')
         assert key == 'a2a:task:test_task'
 
     def test_get_stream_key_empty_task_id(self):
         """Test stream key generation with empty task_id."""
-        injector = RedisStreamInjector()
+        mock_client = AsyncMock()
+        injector = RedisStreamInjector(redis_client=mock_client)
 
         with pytest.raises(ValueError, match='task_id cannot be empty'):
             injector._get_stream_key('')
 
     def test_serialize_event(self):
         """Test event serialization."""
-        injector = RedisStreamInjector()
-
-        data = {'key': 'value', 'number': 42}
-        result = injector._serialize_event('TestEvent', data)
+        injector = RedisStreamInjector(redis_client=AsyncMock())
 
-        assert result['type'] == 'TestEvent'
-        assert 'payload' in result
-
-        # Parse the payload to verify it's correct JSON
-        payload = json.loads(result['payload'])
-        assert payload == data
+        event_data = injector._serialize_event('test_type', {'key': 'value'})
+        assert event_data['type'] == 'test_type'
+        assert 'payload' in event_data
 
     @pytest.mark.asyncio
     async def test_append_to_stream(self):
         """Test appending event to stream."""
         mock_client = AsyncMock()
         mock_client.xadd = AsyncMock(return_value='123-0')
 
-        injector = RedisStreamInjector()
-        injector._client = mock_client
-        injector._connected = True
+        injector = RedisStreamInjector(redis_client=mock_client)
 
         event_data = {'type': 'Test', 'payload': '{"data": "test"}'}
         result = await injector._append_to_stream('test_task', event_data)
@@ -183,7 +183,8 @@ async def test_append_to_stream(self):
     @pytest.mark.asyncio
     async def test_append_to_stream_not_connected(self):
         """Test append_to_stream when not connected."""
-        injector = RedisStreamInjector()
+        mock_client = AsyncMock()
+        injector = RedisStreamInjector(redis_client=mock_client)
         injector._connected = False
 
         with pytest.raises(RuntimeError, match='Not connected to Redis'):
@@ -195,9 +196,7 @@ async def test_stream_message_with_dict(self):
         mock_client = AsyncMock()
         mock_client.xadd = AsyncMock(return_value='123-0')
 
-        injector = RedisStreamInjector()
-        injector._client = mock_client
-        injector._connected = True
+        injector = RedisStreamInjector(redis_client=mock_client)
 
         message_data = {'content': 'test message', 'role': 'assistant'}
         result = await injector.stream_message(