fix: keep streaming alive after the server pong reply (#19)

Co-authored-by: deleteLater <mikcczhang@gmail.com>

Author: himanshu007-creator

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...')

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')) \

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()