Skip to content

Commit bf71208

Browse files
committed
fix mocks for sse tests
1 parent d390223 commit bf71208

2 files changed

Lines changed: 41 additions & 34 deletions

File tree

growthbook/growthbook_client.py

Lines changed: 25 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -187,31 +187,38 @@ async def _start_sse_refresh(self) -> None:
187187
if self._refresh_task is not None: # Already running
188188
return
189189

190-
# SSEClient invokes `on_event` synchronously from a background thread.
191-
async def _handle_sse_event(event_data: Dict[str, Any]) -> None:
192-
try:
193-
event_type = event_data.get("type")
194-
if event_type == "features-updated":
195-
response = await self.load_features_async(
196-
self._api_host, self._client_key, self._decryption_key, self._cache_ttl
197-
)
198-
if response is not None:
199-
await self._handle_feature_update(response)
200-
elif event_type == "features":
201-
data = event_data.get("data", "{}")
202-
if isinstance(data, str):
203-
data = json.loads(data)
204-
await self._handle_feature_update(data)
205-
except Exception:
206-
logger.exception("Error handling SSE event")
190+
async def _handle_sse_event(self, event_data: Dict[str, Any]) -> None:
191+
"""Process an event received from the SSE connection"""
192+
try:
193+
event_type = event_data.get("type")
194+
if event_type == "features-updated":
195+
response = await self.load_features_async(
196+
self._api_host, self._client_key, self._decryption_key, self._cache_ttl
197+
)
198+
if response is not None:
199+
await self._handle_feature_update(response)
200+
elif event_type == "features":
201+
data = event_data.get("data", "{}")
202+
if isinstance(data, str):
203+
data = json.loads(data)
204+
await self._handle_feature_update(data)
205+
except Exception:
206+
logger.exception("Error handling SSE event")
207207

208+
async def _start_sse_refresh(self) -> None:
209+
"""Start SSE-based feature refresh"""
210+
with self._refresh_lock:
211+
if self._refresh_task is not None: # Already running
212+
return
213+
214+
# SSEClient invokes `on_event` synchronously from a background thread.
208215
main_loop = asyncio.get_running_loop()
209216

210217
# We must not pass an `async def` callback here (it would never be awaited).
211218
def sse_handler(event_data: Dict[str, Any]) -> None:
212219
# Schedule async processing onto the main event loop.
213220
try:
214-
asyncio.run_coroutine_threadsafe(_handle_sse_event(event_data), main_loop)
221+
asyncio.run_coroutine_threadsafe(self._handle_sse_event(event_data), main_loop)
215222
except Exception:
216223
logger.exception("Failed to schedule SSE event handler")
217224

tests/test_growthbook_client.py

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -95,8 +95,13 @@ async def test_sse_connection_lifecycle(mock_options, mock_features_response):
9595
patch('growthbook.growthbook_client.EnhancedFeatureRepository.stopAutoRefresh') as mock_stop:
9696
await client.initialize()
9797
# Allow the SSE lifecycle task to start and invoke startAutoRefresh
98-
await asyncio.sleep(0)
98+
await asyncio.sleep(0.1)
9999
assert mock_start.called
100+
101+
# Verify the thread created is a daemon thread (if possible without real start)
102+
# Since we mock startAutoRefresh, we can't check the real thread here.
103+
# But we can check that SSEClient is initialized correctly if we don't mock it all.
104+
100105
await client.close()
101106
assert mock_stop.called
102107

@@ -337,18 +342,14 @@ async def test_callback(features):
337342

338343
@pytest.mark.asyncio
339344
async def test_sse_event_handling(mock_options):
340-
"""Test SSE event handling and reconnection logic"""
345+
"""Test SSE event handling including JSON parsing"""
341346
events = [
342-
{'type': 'features', 'data': {'features': {'feature1': {'defaultValue': 1}}}},
343-
{'type': 'ping', 'data': {}}, # Should be ignored
344-
{'type': 'features', 'data': {'features': {'feature1': {'defaultValue': 2}}}}
347+
# Real SSE payload is a raw string in 'data'
348+
{'type': 'features', 'data': json.dumps({'features': {'feature1': {'defaultValue': 1}}})},
349+
{'type': 'ping', 'data': '{}'}, # Should be ignored
350+
{'type': 'features', 'data': json.dumps({'features': {'feature1': {'defaultValue': 2}}})}
345351
]
346352

347-
async def mock_sse_handler(event_data):
348-
"""Mock the SSE event handler to directly update feature cache"""
349-
if event_data['type'] == 'features':
350-
await client._features_repository._handle_feature_update(event_data['data'])
351-
352353
with patch('growthbook.FeatureRepository.load_features_async',
353354
new_callable=AsyncMock, return_value={"features": {}, "savedGroups": {}}) as mock_load:
354355

@@ -364,16 +365,15 @@ async def mock_sse_handler(event_data):
364365
try:
365366
await client.initialize()
366367

367-
# Simulate SSE events directly
368+
# Simulate SSE events using the actual handler method
369+
# This now tests the json.loads parsing logic!
368370
for event in events:
369-
if event['type'] == 'features':
370-
await client._features_repository._handle_feature_update(event['data'])
371+
await client._features_repository._handle_sse_event(event)
371372

372-
# print(f"AFTER TEST: Current cache state: {client._features_repository._feature_cache.get_current_state()}")
373373
# Verify feature update happened
374-
assert client._features_repository._feature_cache.get_current_state()["features"]["feature1"]["defaultValue"] == 2
374+
state = client._features_repository._feature_cache.get_current_state()
375+
assert state["features"]["feature1"]["defaultValue"] == 2
375376
finally:
376-
# Ensure we clean up the SSE connection
377377
await client.close()
378378

379379
@pytest.mark.asyncio

0 commit comments

Comments
 (0)