diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/extensions/websockets/openbb_websockets/client.py index bfd9c6f0cc1..2c669035d56 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/client.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/client.py @@ -537,6 +537,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes try: with sqlite3.connect(self.results_path) as conn: + conn.execute("PRAGMA journal_mode=WAL;") conn.execute(f"DELETE FROM {self.table_name}") # noqa conn.commit() diff --git a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py index 9c19fb8ae33..62d7c36b9fd 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py @@ -170,8 +170,8 @@ async def setup_database(results_path, table_name): import os # noqa import aiosqlite - async with aiosqlite.connect(results_path) as conn: - if os.path.exists(results_path): + if os.path.exists(results_path): + async with aiosqlite.connect(results_path) as conn: try: await conn.execute("SELECT name FROM sqlite_master WHERE type='table';") except aiosqlite.DatabaseError as e: @@ -180,6 +180,7 @@ async def setup_database(results_path, table_name): "Please check the path, and inspect the file if it exists." + f" -> {e}" ) from e + async with aiosqlite.connect(results_path) as conn: await conn.execute( f""" @@ -198,7 +199,7 @@ async def write_to_db(message, results_path, table_name, limit): import aiosqlite conn = await aiosqlite.connect(results_path) - + await conn.execute("PRAGMA journal_mode=WAL;") try: # Check if the table exists and create it if it doesn't await conn.execute( 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 deleted file mode 100644 index c4434157cff..00000000000 --- a/openbb_platform/providers/fmp/tests/record/object_hash/test_fmp_fetchers/test_fmp_websocket_fetcher.json +++ /dev/null @@ -1 +0,0 @@ -["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 deleted file mode 100644 index 43bd419589e..00000000000 --- a/openbb_platform/providers/fmp/tests/record/screen/test_fmp_fetchers/test_fmp_websocket_fetcher.json +++ /dev/null @@ -1 +0,0 @@ -{"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 35e6c952db8..6386a4c7b3c 100644 --- a/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py +++ b/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py @@ -1,7 +1,9 @@ """Unit tests for FMP provider modules.""" import re +import time from datetime import date +from unittest.mock import MagicMock, patch import pytest from openbb_core.app.service.user_service import UserService @@ -45,6 +47,7 @@ from openbb_fmp.models.executive_compensation import FMPExecutiveCompensationFet from openbb_fmp.models.financial_ratios import FMPFinancialRatiosFetcher from openbb_fmp.models.forward_ebitda_estimates import FMPForwardEbitdaEstimatesFetcher from openbb_fmp.models.forward_eps_estimates import FMPForwardEpsEstimatesFetcher +from openbb_fmp.models.government_trades import FMPGovernmentTradesFetcher from openbb_fmp.models.historical_dividends import FMPHistoricalDividendsFetcher from openbb_fmp.models.historical_employees import FMPHistoricalEmployeesFetcher from openbb_fmp.models.historical_eps import FMPHistoricalEpsFetcher @@ -69,16 +72,96 @@ from openbb_fmp.models.revenue_geographic import FMPRevenueGeographicFetcher from openbb_fmp.models.risk_premium import FMPRiskPremiumFetcher from openbb_fmp.models.share_statistics import FMPShareStatisticsFetcher from openbb_fmp.models.treasury_rates import FMPTreasuryRatesFetcher -from openbb_fmp.models.websocket_connection import FmpWebSocketFetcher +from openbb_fmp.models.websocket_connection import ( + FmpWebSocketConnection, + FmpWebSocketData, + FmpWebSocketFetcher, +) from openbb_fmp.models.world_news import FMPWorldNewsFetcher from openbb_fmp.models.yield_curve import FMPYieldCurveFetcher -from openbb_fmp.models.government_trades import FMPGovernmentTradesFetcher +from openbb_websockets.client import WebSocketClient test_credentials = UserService().default_user_settings.credentials.model_dump( mode="json" ) +MOCK_WEBSOCKET_DATA = [ + { + "date": "2024-12-12T21:13:17.925000-05:00", + "symbol": "BTCUSD", + "exchange": "gdax", + "type": "trade", + "last_price": 99344.20658975664, + "last_size": 1.1701095899999996, + }, + { + "date": "2024-12-12T21:13:19.134000-05:00", + "symbol": "BTCUSD", + "exchange": "gemini", + "type": "trade", + "last_price": 99356.52, + "last_size": 0.01106662, + }, + { + "date": "2024-12-12T21:13:23.017000-05:00", + "symbol": "BTCUSD", + "exchange": "gdax", + "type": "trade", + "last_price": 99359.9798569503, + "last_size": 0.16290354, + }, + { + "date": "2024-12-12T21:13:27.253000-05:00", + "symbol": "BTCUSD", + "exchange": "gemini", + "type": "trade", + "last_price": 99350.54645901076, + "last_size": 0.10452787999999999, + }, + { + "date": "2024-12-12T21:13:28.332000-05:00", + "symbol": "BTCUSD", + "exchange": "gdax", + "type": "trade", + "last_price": 99351.54026716301, + "last_size": 0.11509677, + }, + { + "date": "2024-12-12T21:13:33.842000-05:00", + "symbol": "BTCUSD", + "exchange": "gdax", + "type": "trade", + "last_price": 99353.72826832562, + "last_size": 0.00476533, + }, + { + "date": "2024-12-12T21:13:37.500000-05:00", + "symbol": "BTCUSD", + "exchange": "kraken", + "type": "trade", + "last_price": 99325.1, + "last_size": 0.00302361, + }, + { + "date": "2024-12-12T21:13:38.764000-05:00", + "symbol": "BTCUSD", + "exchange": "gemini", + "type": "trade", + "last_price": 99376.83, + "last_size": 0.00739754, + }, + { + "date": "2024-12-12T21:13:39.310000-05:00", + "symbol": "BTCUSD", + "exchange": "gdax", + "type": "trade", + "last_price": 99363.03061442816, + "last_size": 0.04905618, + }, +] + + def response_filter(response): """Filter the response.""" if "Location" in response["headers"]: @@ -101,6 +184,66 @@ def vcr_config(): } +@pytest.fixture +def mock_websocket_connection(): + """Mock websocket client.""" + + mock_connection = FmpWebSocketConnection( + client=MagicMock( + spec=WebSocketClient( + name="fmp_test", + module="openbb_fmp.utils.websocket_client", + symbol="btcusd", + limit=10, + data_model=FmpWebSocketData, + url="wss://mock.fmp.com/crypto", + api_key="MOCK_API_KEY", + ) + ) + ) + mock_connection.client.is_running = False + mock_results = [] + + def mock_connect(): + mock_connection.client.is_running = True + for data in MOCK_WEBSOCKET_DATA: + mock_results.append(FmpWebSocketData(**data)) + time.sleep(0.1) + + def mock_get_results(): + return mock_results + + mock_connection.client.connect = mock_connect + mock_connection.client.results = mock_get_results + + return mock_connection + + +@pytest.mark.asyncio +async def test_websocket_fetcher( + mock_websocket_connection, credentials=test_credentials +): + """Test websocket fetcher.""" + fetcher = FmpWebSocketFetcher() + params = { + "symbol": "btcusd", + "name": "fmp_test", + "limit": 10, + "asset_type": "crypto", + } + + with patch.object(fetcher, "fetch_data", return_value=mock_websocket_connection): + result = await fetcher.fetch_data(params, credentials) + + # Ensure the client is not running initially + assert not result.client.is_running + assert result.client.results() == [] + result.client.connect() + assert result.client.is_running + assert len(result.client.results()) == len(MOCK_WEBSOCKET_DATA) + assert result.client.results()[0] == FmpWebSocketData(**MOCK_WEBSOCKET_DATA[0]) + + @pytest.mark.record_http def test_fmp_company_filings_fetcher(credentials=test_credentials): """Test FMP company filings fetcher.""" diff --git a/openbb_platform/providers/intrinio/tests/test_intrinio_fetchers.py b/openbb_platform/providers/intrinio/tests/test_intrinio_fetchers.py index a10769baafa..a49177d2c89 100644 --- a/openbb_platform/providers/intrinio/tests/test_intrinio_fetchers.py +++ b/openbb_platform/providers/intrinio/tests/test_intrinio_fetchers.py @@ -1,7 +1,9 @@ """Test Intrinio fetchers.""" +import time from datetime import date from unittest import mock +from unittest.mock import MagicMock, patch import pytest from openbb_core.app.service.user_service import UserService @@ -65,13 +67,178 @@ from openbb_intrinio.models.search_attributes import ( IntrinioSearchAttributesFetcher, ) from openbb_intrinio.models.share_statistics import IntrinioShareStatisticsFetcher +from openbb_intrinio.models.websocket_connection import ( + IntrinioWebSocketConnection, + IntrinioWebSocketData, + IntrinioWebSocketFetcher, +) from openbb_intrinio.models.world_news import IntrinioWorldNewsFetcher +from openbb_websockets.client import WebSocketClient test_credentials = UserService().default_user_settings.credentials.model_dump( mode="json" ) +MOCK_WEBSOCKET_DATA = [ + { + "date": "2024-12-12T18:27:22.943000-05:00", + "symbol": "DJT", + "exchange": "Bats EDGX Exchange", + "type": "trade", + "price": 36.2400016784668, + "size": 155, + "volume": 15968471, + "conditions": "Regular Sale; Form T", + "is_darkpool": None, + "feed": "UTP", + }, + { + "date": "2024-12-12T18:27:22.873000-05:00", + "symbol": "CLF", + "exchange": "NYSE Arca", + "type": "ask", + "price": 11.0, + "size": 4, + "volume": None, + "conditions": None, + "is_darkpool": None, + "feed": "CTA_A", + }, + { + "date": "2024-12-12T18:27:22.873000-05:00", + "symbol": "CLF", + "exchange": "NYSE Arca", + "type": "bid", + "price": 10.960000038146973, + "size": 3, + "volume": None, + "conditions": None, + "is_darkpool": None, + "feed": "CTA_A", + }, + { + "date": "2024-12-12T18:27:22.752000-05:00", + "symbol": "TSLZ", + "exchange": "NYSE Arca", + "type": "ask", + "price": 2.609999895095825, + "size": 223, + "volume": None, + "conditions": None, + "is_darkpool": None, + "feed": "UTP", + }, + { + "date": "2024-12-12T18:27:22.796000-05:00", + "symbol": "QQQM", + "exchange": "NASDAQ (Tape C securities)", + "type": "ask", + "price": 217.67999267578125, + "size": 56, + "volume": None, + "conditions": None, + "is_darkpool": None, + "feed": "UTP", + }, + { + "date": "2024-12-12T18:27:22.796000-05:00", + "symbol": "QQQM", + "exchange": "NASDAQ (Tape C securities)", + "type": "bid", + "price": 217.38999938964844, + "size": 100, + "volume": None, + "conditions": None, + "is_darkpool": None, + "feed": "UTP", + }, + { + "date": "2024-12-12T18:27:22.796000-05:00", + "symbol": "QQQM", + "exchange": "NYSE Arca", + "type": "trade", + "price": 217.5500030517578, + "size": 300, + "volume": 2234058, + "conditions": "Regular Sale; Form T", + "is_darkpool": None, + "feed": "UTP", + }, + { + "date": "2024-12-12T18:27:22.935000-05:00", + "symbol": "LCTX", + "exchange": "Bats BZX Exchange", + "type": "ask", + "price": 0.5356000065803528, + "size": 12, + "volume": None, + "conditions": None, + "is_darkpool": None, + "feed": "CTA_B", + }, + { + "date": "2024-12-12T18:27:22.870000-05:00", + "symbol": "BITU", + "exchange": "Bats EDGX Exchange", + "type": "ask", + "price": 59.5, + "size": 1, + "volume": None, + "conditions": None, + "is_darkpool": None, + "feed": "CTA_B", + }, + { + "date": "2024-12-12T18:27:22.870000-05:00", + "symbol": "BITU", + "exchange": "Bats EDGX Exchange", + "type": "bid", + "price": 58.33000183105469, + "size": 1, + "volume": None, + "conditions": None, + "is_darkpool": None, + "feed": "CTA_B", + }, +] + + +@pytest.fixture +def mock_websocket_connection(): + """Mock websocket client.""" + + mock_connection = IntrinioWebSocketConnection( + client=MagicMock( + spec=WebSocketClient( + name="intrinio_test", + module="openbb_intrinio.utils.websocket_client", + symbol="btcusd", + limit=10, + data_model=IntrinioWebSocketData, + url="wss://mock.intrinio.com/", + api_key="MOCK_API_KEY", + ) + ) + ) + mock_connection.client.is_running = False + mock_results = [] + + def mock_connect(): + mock_connection.client.is_running = True + for data in MOCK_WEBSOCKET_DATA: + mock_results.append(IntrinioWebSocketData(**data)) + time.sleep(0.1) + + def mock_get_results(): + return mock_results + + mock_connection.client.connect = mock_connect + mock_connection.client.results = mock_get_results + + return mock_connection + + @pytest.fixture(scope="module") def vcr_config(): """VCR configuration.""" @@ -99,6 +266,33 @@ def mock_cpu_count(): yield +@pytest.mark.asyncio +async def test_websocket_fetcher( + mock_websocket_connection, credentials=test_credentials +): + """Test websocket fetcher.""" + fetcher = IntrinioWebSocketFetcher() + params = { + "symbol": "*", + "name": "intrinio_test", + "limit": 10, + "asset_type": "stock", + } + + with patch.object(fetcher, "fetch_data", return_value=mock_websocket_connection): + result = await fetcher.fetch_data(params, credentials) + + # Ensure the client is not running initially + assert not result.client.is_running + assert result.client.results() == [] + result.client.connect() + assert result.client.is_running + assert len(result.client.results()) == len(MOCK_WEBSOCKET_DATA) + assert result.client.results()[0] == IntrinioWebSocketData( + **MOCK_WEBSOCKET_DATA[0] + ) + + @pytest.mark.record_http def test_intrinio_equity_historical_fetcher(credentials=test_credentials): """Test equity historical fetcher.""" 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 deleted file mode 100644 index bacd124ff3f..00000000000 --- a/openbb_platform/providers/polygon/tests/record/object_hash/test_polygon_fetchers/test_polygon_websocket_fetcher.json +++ /dev/null @@ -1 +0,0 @@ -["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 deleted file mode 100644 index 0282266fe3e..00000000000 --- a/openbb_platform/providers/polygon/tests/record/screen/test_polygon_fetchers/test_polygon_websocket_fetcher.json +++ /dev/null @@ -1 +0,0 @@ -{"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 cfae156c4e0..e62c458c331 100644 --- a/openbb_platform/providers/polygon/tests/test_polygon_fetchers.py +++ b/openbb_platform/providers/polygon/tests/test_polygon_fetchers.py @@ -1,6 +1,8 @@ """Test the Polygon fetchers.""" +import time from datetime import date +from unittest.mock import MagicMock, patch import pytest from openbb_core.app.service.user_service import UserService @@ -17,13 +19,131 @@ 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 +from openbb_polygon.models.websocket_connection import ( + PolygonWebSocketConnection, + PolygonWebSocketData, + PolygonWebSocketFetcher, +) +from openbb_websockets.client import WebSocketClient test_credentials = UserService().default_user_settings.credentials.model_dump( mode="json" ) +MOCK_WEBSOCKET_DATA = [ + { + "date": "2024-12-12T21:04:02-05:00", + "symbol": "BTC-USD", + "type": "XAS", + "open": 99445.07, + "high": 99445.08, + "low": 99445.07, + "close": 99445.08, + "volume": 0.00188791, + "vwap": 99445.0778, + }, + { + "date": "2024-12-12T21:04:03-05:00", + "symbol": "BTC-USD", + "type": "XAS", + "open": 99445.07, + "high": 99445.08, + "low": 99445.07, + "close": 99445.07, + "volume": 0.00670653, + "vwap": 99445.0702, + }, + { + "date": "2024-12-12T21:04:04-05:00", + "symbol": "BTC-USD", + "type": "XAS", + "open": 99445.07, + "high": 99445.07, + "low": 99445.07, + "close": 99445.07, + "volume": 0.00007158, + "vwap": 99445.07, + }, + { + "date": "2024-12-12T21:04:05-05:00", + "symbol": "BTC-USD", + "type": "XAS", + "open": 99445.07, + "high": 99445.07, + "low": 99428.34, + "close": 99428.34, + "volume": 0.6058888, + "vwap": 99441.6093, + }, + { + "date": "2024-12-12T21:04:06-05:00", + "symbol": "BTC-USD", + "type": "XAS", + "open": 99428.33, + "high": 99428.34, + "low": 99428.33, + "close": 99428.34, + "volume": 0.01953452, + "vwap": 99428.3301, + }, + { + "date": "2024-12-12T21:04:07-05:00", + "symbol": "BTC-USD", + "type": "XAS", + "open": 99428.33, + "high": 99428.34, + "low": 99428.33, + "close": 99428.33, + "volume": 0.1214753, + "vwap": 99428.33, + }, + { + "date": "2024-12-12T21:04:08-05:00", + "symbol": "BTC-USD", + "type": "XAS", + "open": 99428.33, + "high": 99435.48, + "low": 99428.33, + "close": 99435.48, + "volume": 0.20089819, + "vwap": 99429.29, + }, + { + "date": "2024-12-12T21:04:09-05:00", + "symbol": "BTC-USD", + "type": "XAS", + "open": 99435.47, + "high": 99435.47, + "low": 99435.02, + "close": 99435.02, + "volume": 0.03098464, + "vwap": 99435.3318, + }, + { + "date": "2024-12-12T21:04:10-05:00", + "symbol": "BTC-USD", + "type": "XAS", + "open": 99445.08, + "high": 99445.08, + "low": 99445.05, + "close": 99445.05, + "volume": 0.00245657, + "vwap": 99445.0502, + }, + { + "date": "2024-12-12T21:04:11-05:00", + "symbol": "BTC-USD", + "type": "XAS", + "open": 99445.08, + "high": 99445.08, + "low": 99440.55, + "close": 99440.55, + "volume": 0.06000562, + "vwap": 99443.6374, + }, +] + @pytest.fixture(scope="module") def vcr_config(): @@ -36,6 +156,69 @@ def vcr_config(): } +@pytest.fixture +def mock_websocket_connection(): + """Mock websocket client.""" + + mock_connection = PolygonWebSocketConnection( + client=MagicMock( + spec=WebSocketClient( + name="polygon_test", + module="openbb_polygon.utils.websocket_client", + symbol="btcusd", + limit=10, + data_model=PolygonWebSocketData, + url="wss://mock.polygon.com/crypto", + api_key="MOCK_API_KEY", + ) + ) + ) + mock_connection.client.is_running = False + mock_results = [] + + def mock_connect(): + mock_connection.client.is_running = True + for data in MOCK_WEBSOCKET_DATA: + mock_results.append(PolygonWebSocketData(**data)) + time.sleep(0.1) + + def mock_get_results(): + return mock_results + + mock_connection.client.connect = mock_connect + mock_connection.client.results = mock_get_results + + return mock_connection + + +@pytest.mark.asyncio +async def test_websocket_fetcher( + mock_websocket_connection, credentials=test_credentials +): + """Test websocket fetcher.""" + fetcher = PolygonWebSocketFetcher() + params = { + "symbol": "btcusd", + "name": "polygon_test", + "limit": 10, + "asset_type": "crypto", + "feed": "aggs_sec", + } + + with patch.object(fetcher, "fetch_data", return_value=mock_websocket_connection): + result = await fetcher.fetch_data(params, credentials) + + # Ensure the client is not running initially + assert not result.client.is_running + assert result.client.results() == [] + result.client.connect() + assert result.client.is_running + assert len(result.client.results()) == len(MOCK_WEBSOCKET_DATA) + assert result.client.results()[0] == PolygonWebSocketData( + **MOCK_WEBSOCKET_DATA[0] + ) + + @pytest.mark.record_http def test_polygon_equity_historical_fetcher(credentials=test_credentials): """Test the Polygon Equity Historical fetcher.""" @@ -171,33 +354,3 @@ 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 deleted file mode 100644 index d46ccbdbefb..00000000000 --- a/openbb_platform/providers/tiingo/tests/record/object_hash/test_tiingo_fetchers/test_tiingo_websocket_fetcher.json +++ /dev/null @@ -1 +0,0 @@ -["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 deleted file mode 100644 index 6f58aa2e25a..00000000000 --- a/openbb_platform/providers/tiingo/tests/record/screen/test_tiingo_fetchers/test_tiingo_websocket_fetcher.json +++ /dev/null @@ -1 +0,0 @@ -{"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 a08bbf43640..8fbf39aff7b 100644 --- a/openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py +++ b/openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py @@ -1,6 +1,8 @@ """Test Tiingo fetchers.""" +import time from datetime import date +from unittest.mock import MagicMock, patch import pytest from openbb_core.app.service.user_service import UserService @@ -9,13 +11,85 @@ 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.websocket_connection import ( + TiingoWebSocketConnection, + TiingoWebSocketData, + TiingoWebSocketFetcher, +) from openbb_tiingo.models.world_news import TiingoWorldNewsFetcher +from openbb_websockets.client import WebSocketClient test_credentials = UserService().default_user_settings.credentials.model_dump( mode="json" ) +MOCK_WEBSOCKET_DATA = [ + { + "date": "2024-12-12 16:57:50.164993-0500", + "symbol": "BTCUSD", + "type": "trade", + "exchange": "gdax", + "last_price": 99818.08814355676, + "last_size": 0.01445296, + }, + { + "date": "2024-12-12 16:57:50.697317-0500", + "symbol": "BTCUSD", + "type": "trade", + "exchange": "kraken", + "last_price": 99810.1, + "last_size": 6.856e-05, + }, + { + "date": "2024-12-12 16:57:51.119000-0500", + "symbol": "BTCUSD", + "type": "trade", + "exchange": "gemini", + "last_price": 99827.02771283902, + "last_size": 0.19860106, + }, + { + "date": "2024-12-12 16:57:52.573000-0500", + "symbol": "BTCUSD", + "type": "trade", + "exchange": "bitfinex", + "last_price": 99780.0, + "last_size": 8.6e-05, + }, + { + "date": "2024-12-12 16:57:55.187865-0500", + "symbol": "BTCUSD", + "type": "trade", + "exchange": "gdax", + "last_price": 99837.23198173672, + "last_size": 1.50565886, + }, + { + "date": "2024-12-12 16:57:55-0500", + "symbol": "BTCUSD", + "type": "trade", + "exchange": "bitstamp", + "last_price": 99862.0, + "last_size": 0.00212959, + }, + { + "date": "2024-12-12 16:57:57.647609-0500", + "symbol": "BTCUSD", + "type": "trade", + "exchange": "kraken", + "last_price": 99819.24927536234, + "last_size": 0.00207, + }, + { + "date": "2024-12-12 16:58:00.009694-0500", + "symbol": "BTCUSD", + "type": "trade", + "exchange": "gdax", + "last_price": 99838.34191368945, + "last_size": 0.037644090000000005, + }, +] + @pytest.fixture(scope="module") def vcr_config(): @@ -28,6 +102,68 @@ def vcr_config(): } +@pytest.fixture +def mock_websocket_connection(): + """Mock websocket client.""" + + mock_connection = TiingoWebSocketConnection( + client=MagicMock( + spec=WebSocketClient( + name="tiingo_test", + module="openbb_tiingo.utils.websocket_client", + symbol="btcusd", + limit=10, + data_model=TiingoWebSocketData, + url="wss://mock.tiingo.com/iex", + api_key="MOCK_TOKEN", + ) + ) + ) + mock_connection.client.is_running = False + mock_results = [] + + def mock_connect(): + mock_connection.client.is_running = True + for data in MOCK_WEBSOCKET_DATA: + mock_results.append(TiingoWebSocketData(**data)) + time.sleep(0.1) + + def mock_get_results(): + return mock_results + + mock_connection.client.connect = mock_connect + mock_connection.client.results = mock_get_results + + return mock_connection + + +@pytest.mark.asyncio +async def test_websocket_fetcher( + mock_websocket_connection, credentials=test_credentials +): + """Test websocket fetcher.""" + fetcher = TiingoWebSocketFetcher() + params = { + "symbol": "btcusd", + "name": "tiingo_test", + "limit": 10, + "asset_type": "crypto", + } + + with patch.object(fetcher, "fetch_data", return_value=mock_websocket_connection): + result = await fetcher.fetch_data(params, credentials) + + # Ensure the client is not running initially + assert not result.client.is_running + assert result.client.results() == [] + result.client.connect() + assert result.client.is_running + assert len(result.client.results()) == len(MOCK_WEBSOCKET_DATA) + assert result.client.results()[0] == TiingoWebSocketData( + **MOCK_WEBSOCKET_DATA[0] + ) + + @pytest.mark.record_http def test_tiingo_equity_historical_fetcher(credentials=test_credentials): """Test Tiingo equity historical fetcher.""" @@ -98,32 +234,3 @@ 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)