diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/extensions/websockets/openbb_websockets/client.py index 79e2ae2cde7..51f0e39176d 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/client.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/client.py @@ -288,8 +288,10 @@ class WebSocketClient: self._thread.join() err = ChildProcessError(output) self._exception = err - sys.stdout.write(output + "\n") - sys.stdout.flush() + output = output + "\n" + self.logger.info(output) + # sys.stdout.write(output) + # sys.stdout.flush() break output = clean_message(output) @@ -299,9 +301,10 @@ class WebSocketClient: elif output.startswith("INFO:"): output = output.replace("INFO:", "PROVIDER INFO:") - output = output + "\n" - sys.stdout.write(output + "\n") - sys.stdout.flush() + # output = output + "\n" + # sys.stdout.write(output + "\n") + self.logger.info(output) + # sys.stdout.flush() except queue.Empty: continue @@ -534,25 +537,28 @@ class WebSocketClient: # pylint: disable=import-outside-toplevel import json # noqa import sqlite3 + from openbb_core.app.model.abstract.error import OpenBBError output: list = [] file_path = self.results_path if file_path.exists(): with sqlite3.connect(file_path) as conn: - cursor = conn.execute(f"SELECT * FROM {self.table_name}") # noqa - for row in cursor: - index, message = row - if self.data_model: - output.append( - self.data_model.model_validate_json( - json.loads(message) - if isinstance(message, str) - else message - ) - ) - else: - output.append(json.loads(json.loads(message))) - + try: + cursor = conn.execute(f"SELECT * FROM {self.table_name}") # noqa + for row in cursor: + index, message = row + if self.data_model: + message = json.loads(message) + if isinstance(message, (str, bytes)): + output.append( + self.data_model.model_validate_json(message) + ) + elif isinstance(message, dict): + output.append(self.data_model(**message)) + else: + output.append(json.loads(json.loads(message))) + except Exception as e: + raise OpenBBError(f"Error retrieving results: {e}") if output: return output diff --git a/openbb_platform/providers/fmp/tests/record/object_hash/test_fmp_fetchers/test_fmp_websocket_fetcher.json b/openbb_platform/providers/fmp/tests/record/object_hash/test_fmp_fetchers/test_fmp_websocket_fetcher.json new file mode 100644 index 00000000000..c4434157cff --- /dev/null +++ b/openbb_platform/providers/fmp/tests/record/object_hash/test_fmp_fetchers/test_fmp_websocket_fetcher.json @@ -0,0 +1 @@ +["b5bea41b6c623f7c09f1bf24dcae58ebab3c0cdd90ad966bc43a45b44867e12b", "3dd2e537120ffb2eaa71d091dc92478f3ce8aa403bcd17c12fff1c5194a16118", "fcbcf165908dd18a9e49f7ff27810176db8e9f63b4352213741664245224f8aa"] \ No newline at end of file diff --git a/openbb_platform/providers/fmp/tests/record/screen/test_fmp_fetchers/test_fmp_websocket_fetcher.json b/openbb_platform/providers/fmp/tests/record/screen/test_fmp_fetchers/test_fmp_websocket_fetcher.json new file mode 100644 index 00000000000..43bd419589e --- /dev/null +++ b/openbb_platform/providers/fmp/tests/record/screen/test_fmp_fetchers/test_fmp_websocket_fetcher.json @@ -0,0 +1 @@ +{"out": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "err": "PROVIDER INFO: Authenticated\n\nPROVIDER INFO: Subscribed to btcusd\n\nDisconnected from the provider WebSocket.\n\n"} \ No newline at end of file diff --git a/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py b/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py index 563f1b0b414..6ddaa374a7a 100644 --- a/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py +++ b/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py @@ -765,17 +765,25 @@ def test_fmp_historical_market_cap_fetcher(credentials=test_credentials): assert result is None -@pytest.mark.record_screen -def test_fmp_websocket_fetcher(credentials=test_credentials): +@pytest.mark.record_verify_screen(hash=True) +@pytest.mark.record_verify_object(hash=False) +def test_fmp_websocket_fetcher(record, credentials=test_credentials): """Test FMP Websocket fetcher.""" import asyncio + import time params = {"symbol": "btcusd", "name": "test", "limit": 10, "asset_type": "crypto"} try: fetcher = FmpWebSocketFetcher() - result = fetcher.test(params, credentials) response = asyncio.run(fetcher.fetch_data(params, credentials)) - assert result is None + time.sleep(1) + record.add_verify(response.client.is_running) + assert response.client.is_running + time.sleep(1) + assert len(response.client.results) > 0 + record.add_verify(list(response.client.results[0].model_dump().keys())) finally: response.client.disconnect() + assert not response.client.is_running + record.add_verify(response.client.is_running) diff --git a/openbb_platform/providers/polygon/tests/record/object_hash/test_polygon_fetchers/test_polygon_websocket_fetcher.json b/openbb_platform/providers/polygon/tests/record/object_hash/test_polygon_fetchers/test_polygon_websocket_fetcher.json new file mode 100644 index 00000000000..bacd124ff3f --- /dev/null +++ b/openbb_platform/providers/polygon/tests/record/object_hash/test_polygon_fetchers/test_polygon_websocket_fetcher.json @@ -0,0 +1 @@ +["b5bea41b6c623f7c09f1bf24dcae58ebab3c0cdd90ad966bc43a45b44867e12b", "1151296bae1bacb7626f9db0e63520e1bc2a193755e1d88ff58ab1f20abcbaed", "fcbcf165908dd18a9e49f7ff27810176db8e9f63b4352213741664245224f8aa"] \ No newline at end of file diff --git a/openbb_platform/providers/polygon/tests/record/screen/test_polygon_fetchers/test_polygon_websocket_fetcher.json b/openbb_platform/providers/polygon/tests/record/screen/test_polygon_fetchers/test_polygon_websocket_fetcher.json new file mode 100644 index 00000000000..0282266fe3e --- /dev/null +++ b/openbb_platform/providers/polygon/tests/record/screen/test_polygon_fetchers/test_polygon_websocket_fetcher.json @@ -0,0 +1 @@ +{"out": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "err": "PROVIDER INFO: Connected Successfully\n\nPROVIDER INFO: authenticated\n\nPROVIDER INFO: subscribed to: XAS.BTC-USD\n\nDisconnected from the provider WebSocket.\n\n"} \ No newline at end of file diff --git a/openbb_platform/providers/polygon/tests/test_polygon_fetchers.py b/openbb_platform/providers/polygon/tests/test_polygon_fetchers.py index 4e8676b3079..cfae156c4e0 100644 --- a/openbb_platform/providers/polygon/tests/test_polygon_fetchers.py +++ b/openbb_platform/providers/polygon/tests/test_polygon_fetchers.py @@ -17,6 +17,7 @@ from openbb_polygon.models.income_statement import PolygonIncomeStatementFetcher from openbb_polygon.models.index_historical import ( PolygonIndexHistoricalFetcher, ) +from openbb_polygon.models.websocket_connection import PolygonWebSocketFetcher from openbb_polygon.models.market_snapshots import PolygonMarketSnapshotsFetcher test_credentials = UserService().default_user_settings.credentials.model_dump( @@ -170,3 +171,33 @@ def test_polygon_currency_snapshots_fetcher(credentials=test_credentials): fetcher = PolygonCurrencySnapshotsFetcher() result = fetcher.test(params, credentials) assert result is None + + +@pytest.mark.record_verify_screen(hash=True) +@pytest.mark.record_verify_object(hash=False) +def test_polygon_websocket_fetcher(record, credentials=test_credentials): + """Test Polygon Websocket fetcher.""" + import asyncio + import time + + params = { + "symbol": "btcusd", + "name": "polygon_test", + "limit": 10, + "asset_type": "crypto", + "feed": "aggs_sec", + } + + try: + fetcher = PolygonWebSocketFetcher() + response = asyncio.run(fetcher.fetch_data(params, credentials)) + time.sleep(1) + record.add_verify(response.client.is_running) + assert response.client.is_running + time.sleep(1) + assert len(response.client.results) > 0 + record.add_verify(list(response.client.results[0].model_dump().keys())) + finally: + response.client.disconnect() + assert not response.client.is_running + record.add_verify(response.client.is_running) diff --git a/openbb_platform/providers/tiingo/tests/record/object_hash/test_tiingo_fetchers/test_tiingo_websocket_fetcher.json b/openbb_platform/providers/tiingo/tests/record/object_hash/test_tiingo_fetchers/test_tiingo_websocket_fetcher.json new file mode 100644 index 00000000000..d46ccbdbefb --- /dev/null +++ b/openbb_platform/providers/tiingo/tests/record/object_hash/test_tiingo_fetchers/test_tiingo_websocket_fetcher.json @@ -0,0 +1 @@ +["b5bea41b6c623f7c09f1bf24dcae58ebab3c0cdd90ad966bc43a45b44867e12b", "22eb8665ee9e73f51d599865e64d3d03b16cda20434665fa481080c201dddab9", "fcbcf165908dd18a9e49f7ff27810176db8e9f63b4352213741664245224f8aa"] \ No newline at end of file diff --git a/openbb_platform/providers/tiingo/tests/record/screen/test_tiingo_fetchers/test_tiingo_websocket_fetcher.json b/openbb_platform/providers/tiingo/tests/record/screen/test_tiingo_fetchers/test_tiingo_websocket_fetcher.json new file mode 100644 index 00000000000..6f58aa2e25a --- /dev/null +++ b/openbb_platform/providers/tiingo/tests/record/screen/test_tiingo_fetchers/test_tiingo_websocket_fetcher.json @@ -0,0 +1 @@ +{"out": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "err": "PROVIDER INFO: WebSocket connection established.\n\nPROVIDER INFO: Authorization: Success\n\nDisconnected from the provider WebSocket.\n\n"} \ No newline at end of file diff --git a/openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py b/openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py index 26dcddaf1df..a08bbf43640 100644 --- a/openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py +++ b/openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py @@ -9,6 +9,7 @@ from openbb_tiingo.models.crypto_historical import TiingoCryptoHistoricalFetcher from openbb_tiingo.models.currency_historical import TiingoCurrencyHistoricalFetcher from openbb_tiingo.models.equity_historical import TiingoEquityHistoricalFetcher from openbb_tiingo.models.trailing_dividend_yield import TiingoTrailingDivYieldFetcher +from openbb_tiingo.models.websocket_connection import TiingoWebSocketFetcher from openbb_tiingo.models.world_news import TiingoWorldNewsFetcher test_credentials = UserService().default_user_settings.credentials.model_dump( @@ -97,3 +98,32 @@ def test_tiingo_trailing_div_yield_fetcher(credentials=test_credentials): fetcher = TiingoTrailingDivYieldFetcher() result = fetcher.test(params, credentials) assert result is None + + +@pytest.mark.record_verify_screen(hash=True) +@pytest.mark.record_verify_object(hash=False) +def test_tiingo_websocket_fetcher(record, credentials=test_credentials): + """Test Tiingo Websocket fetcher.""" + import asyncio + import time + + params = { + "symbol": "btcusd", + "name": "tiingo_test", + "limit": 10, + "asset_type": "crypto", + } + + try: + fetcher = TiingoWebSocketFetcher() + response = asyncio.run(fetcher.fetch_data(params, credentials)) + time.sleep(1) + record.add_verify(response.client.is_running) + assert response.client.is_running + time.sleep(1) + assert len(response.client.results) > 0 + record.add_verify(list(response.client.results[0].model_dump().keys())) + finally: + response.client.disconnect() + assert not response.client.is_running + record.add_verify(response.client.is_running)