from datetime import UTC from time import sleep, time import pytest from conftest import WAIT_TIME, create_tl_wss_client from tl_shared.models import ( Accounts, AccountState, Orders, QuoteLevel1, ) from tl_shared.types import OrderType, RouteId, SideType, TradableInstrumentId from tl_shared.utils import datetime_tz_aware_now, timestamp_sec_to_datetime_tz_aware from tradelocker.tl_api_client import TLAPIClient from tradelocker_wss.tl_wss_client import TLWSSClient from tradelocker_wss.types import ( MessageType, SocketEndpointType, SocketSubscribeType, ) WAIT_FOR_UPDATE_MS: int = 120_000 # This should be the last timestamp of price_update last_prices_update: int = 0 last_account_update: int = 0 last_orders_update: int = 0 def get_ts() -> int: return int(time() * 1000) # ============== HELPER FUNCTIONS ================== def handle_prices_message(price_update: MessageType) -> None: assert isinstance(price_update, QuoteLevel1) # Modify the global did_receive_prices variable global last_prices_update last_prices_update = get_ts() print(f"@ {last_prices_update} handle_prices_message: {price_update}") def handle_account_message( account_update: MessageType, ) -> None: assert isinstance(account_update, (Accounts, AccountState)) # Modify the global did_receive_prices variable global last_account_update last_account_update = get_ts() print(f"@ {last_account_update} handle_account_message: {account_update}") def handle_orders_message( orders_update: MessageType, ) -> None: assert isinstance(orders_update, Orders) # Modify the global did_receive_prices variable global last_orders_update last_orders_update = get_ts() print(f"@ {last_orders_update} handle_orders_message: {orders_update}") # ============ MAIN TESTS ================= def test_instrument_id(instrument_id_fixture: TradableInstrumentId) -> None: assert instrument_id_fixture def test_subscribe_to_accounts_and_stop( tl_api_fixture: TLAPIClient, tl_wss_client_fixture: TLWSSClient, instrument_id_fixture: int, ) -> None: tl_wss_client = tl_wss_client_fixture assert not tl_wss_client.is_connected(SocketEndpointType.SOCKET_ACCOUNTS) time_start = get_ts() assert last_account_update < time_start tl_wss_client.create_and_start_endpoint( endpoint_type=SocketEndpointType.SOCKET_ACCOUNTS, message_callback=handle_account_message, ) sleep(WAIT_TIME) order_id = tl_api_fixture.trading.place_order( type_=OrderType.MARKET, instrument_id=instrument_id_fixture, qty=0.01, side=SideType.BUY, ) while get_ts() - time_start < WAIT_FOR_UPDATE_MS: if last_account_update > time_start and tl_wss_client.is_connected( SocketEndpointType.SOCKET_ACCOUNTS ): break sleep(WAIT_TIME) assert tl_wss_client.is_connected(SocketEndpointType.SOCKET_ACCOUNTS) assert last_account_update > time_start tl_wss_client.stop() sleep(WAIT_TIME) # Give it some time to stop time_stop_plus = get_ts() assert not tl_wss_client.is_connected(SocketEndpointType.SOCKET_ACCOUNTS) assert last_account_update < time_stop_plus # Close the position opened above tl_api_fixture.trading.close_position(order_id=order_id) def test_subscribe_to_prices_and_stop( tl_wss_client_fixture: TLWSSClient, instrument_id_fixture: TradableInstrumentId, route_id_fixture: RouteId, ) -> None: tl_wss_client = tl_wss_client_fixture assert not tl_wss_client.is_connected(SocketEndpointType.SOCKET_QUOTES) time_start = get_ts() assert last_prices_update < time_start tl_wss_client.create_and_start_endpoint( endpoint_type=SocketEndpointType.SOCKET_QUOTES, message_callback=handle_prices_message, ) sleep(WAIT_TIME) assert tl_wss_client.is_started(SocketEndpointType.SOCKET_QUOTES) # this seems to take some time to start tl_wss_client.subscribe_to_quote( SocketSubscribeType.LEVEL1, instrument_id=instrument_id_fixture, route_id=route_id_fixture, ) sleep(WAIT_TIME) while get_ts() - time_start < WAIT_FOR_UPDATE_MS: if last_prices_update > time_start and tl_wss_client.is_connected( SocketEndpointType.SOCKET_QUOTES ): break sleep(WAIT_TIME) assert tl_wss_client.is_connected(SocketEndpointType.SOCKET_QUOTES) assert last_prices_update > time_start tl_wss_client.stop() sleep(WAIT_TIME) # Give it some time to stop time_stop_plus = get_ts() assert not tl_wss_client.is_connected(SocketEndpointType.SOCKET_QUOTES) assert last_prices_update < time_stop_plus def test_subscribe_to_trade_events_and_stop( tl_api_fixture: TLAPIClient, tl_wss_client_fixture: TLWSSClient, instrument_id_fixture: TradableInstrumentId, ) -> None: tl_api = tl_api_fixture tl_wss_client = tl_wss_client_fixture assert not tl_wss_client.is_connected(SocketEndpointType.SOCKET_TRADES) time_start = get_ts() assert last_orders_update < time_start tl_wss_client.create_and_start_endpoint( endpoint_type=SocketEndpointType.SOCKET_TRADES, message_callback=handle_orders_message, ) sleep(WAIT_TIME) order_id = tl_api.trading.place_order( type_=OrderType.MARKET, instrument_id=instrument_id_fixture, qty=0.01, side=SideType.BUY, ) while get_ts() - time_start < WAIT_FOR_UPDATE_MS: if last_orders_update > time_start and tl_wss_client.is_connected( SocketEndpointType.SOCKET_TRADES ): break sleep(WAIT_TIME) assert tl_wss_client.is_connected(SocketEndpointType.SOCKET_TRADES) assert last_orders_update > time_start tl_wss_client.stop() sleep(WAIT_TIME) # Give it some time to stop time_stop_plus = get_ts() assert not tl_wss_client.is_connected(SocketEndpointType.SOCKET_TRADES) assert last_orders_update < time_stop_plus # Close the position opened above tl_api.trading.close_position(order_id=order_id) def test_auth_token(tl_api_fixture: TLAPIClient) -> None: tl_wss_client = create_tl_wss_client(tl_api_fixture, reconnect=True) tl_wss_client._fetch_new_wss_token() assert tl_wss_client.wss_token is not None assert tl_wss_client.wss_token.payload.exp is not None expiration_dt = timestamp_sec_to_datetime_tz_aware(tl_wss_client.wss_token.payload.exp) assert expiration_dt.tzinfo is not None # check tzinfo exists assert expiration_dt.tzinfo == UTC # check timezone is UTC assert datetime_tz_aware_now() < expiration_dt # check expiration is in future tl_wss_client.stop() if __name__ == "__main__": pytest.main([__file__])