diff --git a/fbclient/streaming.py b/fbclient/streaming.py index 570eecf..443ae48 100644 --- a/fbclient/streaming.py +++ b/fbclient/streaming.py @@ -11,7 +11,7 @@ from fbclient.interfaces import UpdateProcessor from fbclient.notice_broadcaster import NoticeBroadcater from fbclient.status import DataUpdateStatusProviderImpl -from fbclient.status_types import (DATA_INVALID_ERROR, NETWORK_ERROR, +from fbclient.status_types import (NETWORK_ERROR, REQUEST_INVALID_ERROR, RUNTIME_ERROR, SYSTEM_QUIT, UNKNOWN_CLOSE_CODE, UNKNOWN_ERROR, WEBSOCKET_ERROR, State) @@ -302,16 +302,21 @@ def _on_message(self, wsapp: websocket.WebSocketApp, msg): log.trace('Streaming WebSocket data: %s' % msg) # type: ignore try: all_data = json.loads(msg) - if not valide_all_data(all_data): - raise ValueError('invalid streaming data') - if not self._on_process_data(all_data['data']) and self.__wsapp: - # state already updated in init or upsert, just reconn - self.__self_closed = _SelfClosed(is_self_close=True, is_reconn=True, state=None) - wsapp.close(status=WS_GOING_AWAY_CLOSE) - except Exception as e: - self.__self_closed = _SelfClosed(is_self_close=True, is_reconn=False, - state=State.error_off_state(DATA_INVALID_ERROR, str(e))) - wsapp.close(status=WS_GOING_AWAY_CLOSE) + except Exception: + log.exception('FB Python SDK: could not parse data-sync message; skipping message: %r', msg) + return + if isinstance(all_data, dict) and all_data.get('messageType') == 'data-sync': + try: + self._handle_data_sync_message(wsapp, all_data) + except Exception: + log.exception('FB Python SDK: could not process data-sync message; skipping message: %r', all_data) + + def _handle_data_sync_message(self, wsapp: websocket.WebSocketApp, all_data): + if not valide_all_data(all_data): + log.error('FB Python SDK: invalid data-sync message; skipping message: %r', all_data) + return + if not self._on_process_data(all_data['data']): + log.error('FB Python SDK: could not apply data-sync message; skipping message: %r', all_data) def stop(self): log.info('FB Python SDK: Streaming is stopping...') diff --git a/fbclient/utils/__init__.py b/fbclient/utils/__init__.py index 51abb9e..d058e94 100644 --- a/fbclient/utils/__init__.py +++ b/fbclient/utils/__init__.py @@ -67,7 +67,7 @@ def unix_timestamp_in_milliseconds(): def valide_all_data(all_data=None) -> bool: return isinstance(all_data, dict) \ - and all_data.get('messageType', 'pong') == 'data-sync' \ + and all_data.get('messageType') == 'data-sync' \ and 'data' in all_data and isinstance(all_data['data'], dict) \ and all(k in all_data['data'] for k in ('eventType', 'featureFlags', 'segments')) \ and any(k == all_data['data']['eventType'] for k in ('full', 'patch')) \ diff --git a/tests/test_runtime_safety.py b/tests/test_runtime_safety.py index 4d4e3d6..53b08e2 100644 --- a/tests/test_runtime_safety.py +++ b/tests/test_runtime_safety.py @@ -5,6 +5,8 @@ from concurrent.futures import ThreadPoolExecutor from pathlib import Path +import pytest + from fbclient.client import FBClient from fbclient.config import Config, HTTPConfig from fbclient.data_storage import InMemoryDataStorage @@ -470,3 +472,88 @@ def close(self, status=None): streaming.stop() broadcaster.stop() streaming_module._WEBSOCKET_TIMEOUT_LOCK.release() + + +class _RecordingWebSocketApp: + def __init__(self): + self.close_statuses = [] + + def close(self, status=None): + self.close_statuses.append(status) + + +def _idle_streaming(): + broadcaster = NoticeBroadcater() + streaming = Streaming(Config(FAKE_ENV_SECRET, FAKE_URL, FAKE_URL), broadcaster, + DataUpdateStatusProviderImpl(InMemoryDataStorage()), threading.Event()) + return streaming, broadcaster + + +@pytest.mark.parametrize('message', [ + '{"messageType":"pong","data":{}}', + '{"messageType":"unknown","data":{}}', + '{"data":{}}', + '[]', + 'null', +]) +def test_streaming_non_data_sync_message_keeps_connection_open(message): + streaming, broadcaster = _idle_streaming() + wsapp = _RecordingWebSocketApp() + try: + streaming._on_message(wsapp, message) + assert wsapp.close_statuses == [] + assert not streaming._Streaming__self_closed() + finally: + streaming.stop() + broadcaster.stop() + + +@pytest.mark.parametrize('message', [ + 'invalid json', + '{"messageType":"data-sync","data":{}}', + '{"messageType":"data-sync","data":null}', +]) +def test_streaming_invalid_message_is_logged_and_next_sync_is_processed(message, mocker): + streaming, broadcaster = _idle_streaming() + wsapp = _RecordingWebSocketApp() + process_data = mocker.patch.object(streaming, '_on_process_data', return_value=True) + error_log = mocker.patch.object(streaming_module.log, 'error') + exception_log = mocker.patch.object(streaming_module.log, 'exception') + try: + streaming._on_message(wsapp, message) + assert wsapp.close_statuses == [] + assert not streaming._Streaming__self_closed() + process_data.assert_not_called() + assert error_log.called or exception_log.called + + streaming._on_message(wsapp, '{"messageType":"data-sync","data":' + '{"eventType":"full","featureFlags":[],"segments":[]}}') + process_data.assert_called_once_with({'eventType': 'full', 'featureFlags': [], 'segments': []}) + assert wsapp.close_statuses == [] + finally: + streaming.stop() + broadcaster.stop() + + +def test_streaming_data_processing_failure_keeps_connection_open(mocker): + streaming, broadcaster = _idle_streaming() + wsapp = _RecordingWebSocketApp() + streaming._Streaming__wsapp = wsapp + process_data = mocker.patch.object(streaming, '_on_process_data', side_effect=[False, True]) + error_log = mocker.patch.object(streaming_module.log, 'error') + message = ('{"messageType":"data-sync","data":' + '{"eventType":"full","featureFlags":[],"segments":[]}}') + try: + streaming._on_message(wsapp, message) + assert wsapp.close_statuses == [] + assert not streaming._Streaming__self_closed() + error_log.assert_called_once() + + streaming._on_message(wsapp, message) + assert process_data.call_count == 2 + assert wsapp.close_statuses == [] + assert not streaming._Streaming__self_closed() + error_log.assert_called_once() + finally: + streaming.stop() + broadcaster.stop()