diff --git a/intriniorealtime/_websocket.py b/intriniorealtime/_websocket.py new file mode 100644 index 0000000..d7a7d40 --- /dev/null +++ b/intriniorealtime/_websocket.py @@ -0,0 +1,39 @@ +import threading +import time +from typing import Optional + + +def close_websocket_app(app) -> None: + """Best-effort close using websocket-client's public API only.""" + if app is None: + return + try: + app.close() + except Exception: + pass + + +def should_abort_handshake( + handshake_event: threading.Event, + stop_flag: Optional[threading.Event], + generation: int, + current_generation, + deadline: float, +) -> bool: + while True: + if handshake_event.is_set(): + return False + if stop_flag is not None and stop_flag.is_set(): + return False + current = current_generation() if callable(current_generation) else current_generation + if generation != current: + return False + remaining = deadline - time.monotonic() + if remaining <= 0: + return True + handshake_event.wait(min(remaining, 0.1)) + + +def _redact_api_key(message, api_key): + text = str(message) + return text.replace(api_key, "") if api_key else text diff --git a/intriniorealtime/equities_client.py b/intriniorealtime/equities_client.py index 7e3b30e..8049f7e 100644 --- a/intriniorealtime/equities_client.py +++ b/intriniorealtime/equities_client.py @@ -10,7 +10,11 @@ from enum import IntEnum, unique from typing import Optional, Dict, Any +from ._websocket import close_websocket_app, should_abort_handshake, _redact_api_key + SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] +CONNECT_TIMEOUT_SECONDS = 30 +STABLE_SESSION_SECONDS = 60.0 _EMPTY_STRING = "" _NAN = float("NAN") REALTIME = "REALTIME" @@ -33,7 +37,7 @@ HEADER_MESSAGE_FORMAT_KEY = "UseNewEquitiesFormat" HEADER_MESSAGE_FORMAT_VALUE = "v2" HEADER_CLIENT_INFORMATION_KEY = "Client-Information" -HEADER_CLIENT_INFORMATION_VALUE = "IntrinioPythonSDKv6.3.0" +HEADER_CLIENT_INFORMATION_VALUE = "IntrinioPythonSDKv6.4.0" class EquitiesQuote: @@ -130,9 +134,10 @@ def __init__(self, configuration: Dict[str, Any], on_trade: Optional[callable], self.logger.addHandler(log_handler) if 'max_queue_size' in configuration: - self.quotes = queue.Queue(maxsize=configuration['max_queue_size']) + self.max_queue_size = configuration['max_queue_size'] else: - self.quotes = queue.Queue(maxsize=MAX_QUEUE_SIZE) + self.max_queue_size = MAX_QUEUE_SIZE + self.quotes = queue.Queue(maxsize=self.max_queue_size) if self.api_key: if not self.valid_api_key(self.api_key): @@ -170,6 +175,13 @@ def __init__(self, configuration: Dict[str, Any], on_trade: Optional[callable], self.joined_channels = set() self.last_queue_warning_time = 0 self.last_self_heal_backoff = -1 + self._stop_event = threading.Event() + self._connect_lock = threading.Lock() + self._disconnect_lock = threading.Lock() + self._lifecycle_lock = threading.RLock() + self._channels_lock = threading.RLock() + self._lifecycle_generation = 0 + self._disconnecting = False self.quote_handler.start() def auth_url(self) -> str: @@ -227,41 +239,160 @@ def do_backoff(self): self.last_self_heal_backoff += 1 i = min(self.last_self_heal_backoff, len(SELF_HEAL_BACKOFFS) - 1) backoff = SELF_HEAL_BACKOFFS[i] - time.sleep(backoff) + self._stop_event.wait(timeout=backoff) def connect(self): - connected = False - while not connected: - try: - self.logger.info("Connecting...") - self.ready = False - self.joined_channels = set() - - if self.ws: - self.ws.close() - time.sleep(3) - - self.refresh_token() - self.refresh_websocket() - connected = True - except Exception as e: - self.logger.error(f"Cannot connect: {repr(e)}") - self.do_backoff() + with self._lifecycle_lock: + if self._disconnecting: + raise RuntimeError("Cannot connect while client is disconnecting") + requested_generation = self._lifecycle_generation + with self._connect_lock: + with self._lifecycle_lock: + if self._disconnecting: + raise RuntimeError("Cannot connect while client is disconnecting") + if requested_generation != self._lifecycle_generation: + return + self._stop_event.clear() + if self.quote_receiver is not None and self.quote_receiver.is_alive() and self.quote_receiver.enabled: + return + + self.last_self_heal_backoff = -1 + while not self._stop_event.is_set(): + try: + self.logger.info("Connecting...") + with self._channels_lock: + self.ready = False + self.joined_channels = set() + self.refresh_quote_handler() + self.refresh_token() + with self._lifecycle_lock: + if (requested_generation != self._lifecycle_generation + or self._stop_event.is_set()): + return + self.refresh_websocket(requested_generation) + return + except Exception as e: + self.logger.error("Cannot connect: %s", _redact_api_key(repr(e), self.api_key)) + self.do_backoff() + + def _finish_incomplete_disconnect(self, receiver, handler, survivors): + for thread in survivors: + thread.join() + with self._lifecycle_lock: + if (receiver is not None + and self.quote_receiver is receiver + and not receiver.is_alive()): + self.quote_receiver = None + self.ws = None + if (handler is not None + and self.quote_handler is handler + and not handler.is_alive()): + self.quote_handler = None + self._disconnecting = False def disconnect(self): - self.ready = False - self.joined_channels = set() - - if self.ws: - self.ws.close() - time.sleep(1) + with self._disconnect_lock: + with self._lifecycle_lock: + if self._disconnecting: + raise RuntimeError("Client disconnect is already in progress") + self._disconnecting = True + self._lifecycle_generation += 1 + self._stop_event.set() + handler = self.quote_handler + receiver = self.quote_receiver + websocket_app = self.ws + if receiver is not None: + receiver.enabled = False + shutdown_deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS + 1 + deferred_completion = False + try: + if handler is not None: + handler.stop() + + close_websocket_app(websocket_app) + + with self._channels_lock: + self.ready = False + self.joined_channels = set() + + if receiver is not None: + if (getattr(receiver, "ident", True) is not None + and receiver is not threading.current_thread()): + receiver.join(timeout=max( + 0.0, shutdown_deadline - time.monotonic())) + if receiver.is_alive(): + close_websocket_app(websocket_app) + else: + finished_websocket = None + with self._lifecycle_lock: + if self.quote_receiver is receiver: + finished_websocket = self.ws + self.quote_receiver = None + self.ws = None + close_websocket_app(finished_websocket) + else: + with self._lifecycle_lock: + if self.ws is websocket_app: + self.ws = None + + if ( + handler is not None + and getattr(handler, "ident", True) is not None + and handler is not threading.current_thread() + ): + handler.join(timeout=max( + 0.0, shutdown_deadline - time.monotonic())) + if handler is not None and not handler.is_alive(): + with self._lifecycle_lock: + if self.quote_handler is handler: + self.quote_handler = None + + current_thread = threading.current_thread() + non_current_survivors = [ + thread for thread in (receiver, handler) + if (thread is not None + and thread is not current_thread + and thread.is_alive()) + ] + if non_current_survivors: + deferred_survivors = list(non_current_survivors) + if (current_thread in (receiver, handler) + and current_thread.is_alive()): + deferred_survivors.append(current_thread) + deferred_completion = True + threading.Thread( + target=self._finish_incomplete_disconnect, + args=(receiver, handler, deferred_survivors), + daemon=True, + ).start() + raise RuntimeError( + "Client disconnect incomplete: threads did not exit") + finally: + if not deferred_completion: + with self._lifecycle_lock: + self._disconnecting = False + + def refresh_quote_handler(self): + handler = self.quote_handler + if handler is not None: + if handler.is_alive() and not handler.stopped: + return + current_thread = threading.current_thread() + replacing_current_stopped_handler = handler is current_thread and handler.stopped + if handler.ident is not None and handler is not current_thread: + handler.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) + if handler.is_alive() and not replacing_current_stopped_handler: + raise RuntimeError("Previous quote handler did not exit") + self.quotes = queue.Queue(maxsize=self.max_queue_size) + self.quote_handler = EquitiesQuoteHandler(self, self.bypass_parsing) + self.quote_handler.start() def refresh_token(self): headers = {HEADER_CLIENT_INFORMATION_KEY: HEADER_CLIENT_INFORMATION_VALUE} if self.api_key: - response = requests.get(self.auth_url(), headers=headers) + response = requests.get(self.auth_url(), headers=headers, timeout=CONNECT_TIMEOUT_SECONDS) else: - response = requests.get(self.auth_url(), auth=(self.username, self.password), headers=headers) + response = requests.get(self.auth_url(), auth=(self.username, self.password), headers=headers, timeout=CONNECT_TIMEOUT_SECONDS) if response.status_code != 200: raise RuntimeError("Auth failed") @@ -269,13 +400,48 @@ def refresh_token(self): self.token = response.text self.logger.info("Authentication successful!") - def refresh_websocket(self): - self.quote_receiver = EquitiesQuoteReceiver(self) - self.quote_receiver.start() - - def on_connect(self): - self.ready = True - self.last_self_heal_backoff = -1 + def refresh_websocket(self, expected_generation=None): + with self._lifecycle_lock: + if (expected_generation is not None + and expected_generation != self._lifecycle_generation): + return + receiver = self.quote_receiver + if receiver is not None: + if receiver.is_alive(): + if receiver.enabled: + return + close_websocket_app(self.ws) + receiver.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) + if receiver.is_alive(): + close_websocket_app(self.ws) + raise RuntimeError("Previous quote receiver did not exit") + with self._lifecycle_lock: + if self.quote_receiver is receiver: + self.quote_receiver = None + self.ws = None + new_receiver = EquitiesQuoteReceiver(self) + with self._lifecycle_lock: + if ((expected_generation is not None + and expected_generation != self._lifecycle_generation) + or self._stop_event.is_set()): + return + self.quote_receiver = new_receiver + try: + new_receiver.start() + except Exception: + if self.quote_receiver is new_receiver: + self.quote_receiver = None + raise + + def on_connect(self, ws=None, receiver=None): + with self._lifecycle_lock: + if (self._stop_event.is_set() + or (receiver is not None and self.quote_receiver is not receiver) + or (ws is not None and self.ws is not ws)): + return + with self._channels_lock: + self.ready = True + self.joined_channels = set() self.refresh_channels() def on_queue_full(self): @@ -287,42 +453,47 @@ def join(self, channels: list[str]): if isinstance(channels, str): channels = [channels] - self.channels = self.channels | set(channels) - self.refresh_channels() + with self._channels_lock: + self.channels = self.channels | set(channels) + self.refresh_channels() def leave(self, channels: list[str]): if isinstance(channels, str): channels = [channels] - self.channels = self.channels - set(channels) - self.refresh_channels() + with self._channels_lock: + self.channels = self.channels - set(channels) + self.refresh_channels() def leave_all(self): - self.channels = set() - self.refresh_channels() + with self._channels_lock: + self.channels = set() + self.refresh_channels() def refresh_channels(self): - if self.ready != True: - return - - # Join new channels - new_channels = self.channels - self.joined_channels - self.logger.debug(f"New channels: {new_channels}") - for channel in new_channels: - msg = self.join_binary_message(channel) - self.ws.send(msg, websocket.ABNF.OPCODE_BINARY) - self.logger.info(f"Joined channel {channel}") - - # Leave old channels - old_channels = self.joined_channels - self.channels - self.logger.debug(f"Old channels: {old_channels}") - for channel in old_channels: - msg = self.leave_binary_message(channel) - self.ws.send(msg, websocket.ABNF.OPCODE_BINARY) - self.logger.info(f"Left channel {channel}") - - self.joined_channels = self.channels.copy() - self.logger.debug(f"Current channels: {self.joined_channels}") + with self._channels_lock: + if self.ready != True or self._stop_event.is_set(): + return + + # Join new channels + new_channels = self.channels - self.joined_channels + self.logger.debug(f"New channels: {new_channels}") + for channel in new_channels: + msg = self.join_binary_message(channel) + self.ws.send(msg, websocket.ABNF.OPCODE_BINARY) + self.joined_channels.add(channel) + self.logger.info(f"Joined channel {channel}") + + # Leave old channels + old_channels = self.joined_channels - self.channels + self.logger.debug(f"Old channels: {old_channels}") + for channel in old_channels: + msg = self.leave_binary_message(channel) + self.ws.send(msg, websocket.ABNF.OPCODE_BINARY) + self.joined_channels.discard(channel) + self.logger.info(f"Left channel {channel}") + + self.logger.debug(f"Current channels: {self.joined_channels}") def join_binary_message(self, channel: str): if channel == "lobby": @@ -364,27 +535,128 @@ def __init__(self, client): self.daemon = True self.client = client self.enabled = True - self.continuation_queue = queue.Queue(100) - self.continuation_lock: threading.Lock = threading.Lock() def run(self): - self.client.ws = websocket.WebSocketApp( - self.client.websocket_url(), - header={HEADER_MESSAGE_FORMAT_KEY: HEADER_MESSAGE_FORMAT_VALUE, HEADER_CLIENT_INFORMATION_KEY: HEADER_CLIENT_INFORMATION_VALUE}, - on_open=self.on_open, - on_close=self.on_close, - on_message=self.on_message, - on_cont_message=self.on_cont_message, - on_error=self.on_error - ) - - self.client.logger.debug("QuoteReceiver ready") - self.client.ws.run_forever(skip_utf8_validation=True) # skip_utf8_validation for more performance + generation = 0 + generation_lock = threading.Lock() + while self.enabled and not self.client._stop_event.is_set(): + with generation_lock: + generation += 1 + handshake_event = threading.Event() + attempt_cancelled = threading.Event() + connect_deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS + session = {"opened": False, "opened_at": 0.0} + app = None + + def on_open_with_timeout( + ws, + *args, + attempt_handshake=handshake_event, + cancelled=attempt_cancelled, + attempt_session=session, + attempt_generation=generation, + deadline=connect_deadline, + ): + rejected = False + with generation_lock: + if (cancelled.is_set() + or attempt_generation != generation + or not self.enabled + or self.client._stop_event.is_set() + or time.monotonic() >= deadline): + cancelled.set() + rejected = True + attempt_handshake.set() + if not rejected: + attempt_session["opened"] = True + attempt_session["opened_at"] = time.monotonic() + if rejected: + close_websocket_app(app) + return + self.on_open(ws) + + # Leave on_cont_message unset so websocket-client aggregates fragmented messages. + app = websocket.WebSocketApp( + self.client.websocket_url(), + header={HEADER_MESSAGE_FORMAT_KEY: HEADER_MESSAGE_FORMAT_VALUE, HEADER_CLIENT_INFORMATION_KEY: HEADER_CLIENT_INFORMATION_VALUE}, + on_open=on_open_with_timeout, + on_close=self.on_close, + on_message=self.on_message, + on_error=self.on_error + ) + with self.client._lifecycle_lock: + if not self.enabled or self.client._stop_event.is_set(): + close_websocket_app(app) + break + self.client.ws = app + + def watchdog( + this_generation=generation, + handshake_event=handshake_event, + attempt_app=app, + cancelled=attempt_cancelled, + deadline=connect_deadline, + ): + if should_abort_handshake( + handshake_event, + self.client._stop_event, + this_generation, + lambda: generation, + deadline, + ): + with generation_lock: + if this_generation != generation: + return + self.client.logger.warning(f"Websocket connect timed out after {CONNECT_TIMEOUT_SECONDS}s") + cancelled.set() + close_websocket_app(attempt_app) + + self.client.logger.debug("QuoteReceiver ready") + threading.Thread(target=watchdog, daemon=True).start() + if not self.enabled or self.client._stop_event.is_set(): + close_websocket_app(app) + break + try: + app.run_forever( + skip_utf8_validation=True, + ) # skip_utf8_validation for more performance + except Exception as e: + self.client.logger.error(f"Websocket ERROR: {repr(e)}") + finally: + # The attempt is over even if the websocket never opened. Wake + # its watchdog so it cannot fire later during reconnect backoff. + handshake_event.set() + + with self.client._channels_lock: + self.client.ready = False + if not self.enabled or self.client._stop_event.is_set(): + break + + self.client.logger.info("Websocket closed, reconnecting...") + skip_wait = ( + session["opened"] + and (time.monotonic() - session["opened_at"]) >= STABLE_SESSION_SECONDS + ) + if skip_wait: + self.client.last_self_heal_backoff = -1 + while self.enabled and not self.client._stop_event.is_set(): + if skip_wait: + skip_wait = False + else: + self.client.do_backoff() + if not self.enabled or self.client._stop_event.is_set(): + break + try: + self.client.refresh_token() + break + except Exception as e: + self.client.logger.error("Cannot connect: %s", _redact_api_key(repr(e), self.client.api_key)) + self.client.logger.debug("QuoteReceiver exiting") def on_open(self, ws): self.client.logger.info("Websocket opened!") - self.client.on_connect() + self.client.on_connect(ws, self) def on_close(self, ws, code, message): self.client.logger.info("Websocket closed!") @@ -392,51 +664,12 @@ def on_close(self, ws, code, message): def on_error(self, ws, error, *args): try: self.client.logger.error(f"Websocket ERROR: {error}") - self.client.connect() - except Exception as e: - self.client.logger.error(f"Error in on_error handler: {repr(e)}; {repr(error)}") - raise e - - def stitch(self): - full = None - while not self.continuation_queue.empty(): - partial = self.continuation_queue.get(True, 1) - if full is None: - full = partial - else: - full = full.join(partial) - return full - - def on_cont_message(self, partial_message, is_last): # The 3rd argument is continue flag. if 0, the data continue - try: - if DEBUGGING: # This is here for performance reasons so we don't use slow reflection on every message. - if isinstance(partial_message, str): - self.client.logger.debug(f"Received partial message (str): {partial_message.encode('utf-8').hex()}") - else: - if isinstance(partial_message, bytes): - self.client.logger.debug(f"Received partial message (hex): {partial_message.hex()}") - #self.client.logger.debug(f"Received partial message (hex): {partial_message.hex()}") - - self.continuation_lock.acquire() try: - if is_last == 0: - self.continuation_queue.put(partial_message) - else: - self.continuation_queue.put(partial_message) - full_message = self.stitch() - self.on_message(self.client.ws, full_message) - finally: - self.continuation_lock.release() - except queue.Full: - self.client.on_queue_full() + ws.close() + except Exception: + pass except Exception as e: - hex_message = "" - if isinstance(partial_message, str): - hex_message = partial_message.encode('utf-8').hex() - else: - if isinstance(partial_message, bytes): - hex_message = partial_message.hex() - self.client.logger.error(f"Websocket on_message ERROR. Message as hex: {hex_message}; error: {repr(e)}") + self.client.logger.error(f"Error in on_error handler: {repr(e)}; {repr(error)}") raise e def on_message(self, ws, message): @@ -467,6 +700,7 @@ def __init__(self, client, bypass_parsing: bool): self.daemon = True self.client = client self.bypass_parsing = bypass_parsing + self._stop_event = threading.Event() self.subprovider_codes = { 0: NO_SUBPROVIDER, 1: CTA_A, @@ -479,6 +713,17 @@ def __init__(self, client, bypass_parsing: bool): 8: EQUITIES_EDGE } + @property + def stopped(self) -> bool: + return self._stop_event.is_set() + + def stop(self): + self._stop_event.set() + try: + self.client.quotes.put_nowait(None) + except queue.Full: + pass + def parse_quote(self, quote_bytes: bytes, start_index: int = 0) -> EquitiesQuote: buffer = memoryview(quote_bytes) symbol_length = buffer[start_index + 2] @@ -523,7 +768,7 @@ def parse_message(self, message_bytes: bytes, start_index: int, backlog_len: int if callable(self.client.on_trade): try: if self.bypass_parsing: - item = message_bytes[start_index:new_start_index - 1] + item = message_bytes[start_index:new_start_index] else: item = self.parse_trade(message_bytes, start_index) self.client.on_trade(item, backlog_len) @@ -533,7 +778,7 @@ def parse_message(self, message_bytes: bytes, start_index: int, backlog_len: int if callable(self.client.on_quote): try: if self.bypass_parsing: - item = message_bytes[start_index:new_start_index - 1] + item = message_bytes[start_index:new_start_index] else: item = self.parse_quote(message_bytes, start_index) self.client.on_quote(item, backlog_len) @@ -544,11 +789,27 @@ def parse_message(self, message_bytes: bytes, start_index: int, backlog_len: int def run(self): self.client.logger.debug("QuoteHandler ready") - while True: + while not self._stop_event.is_set(): message = self.client.quotes.get() + if message is None or self._stop_event.is_set(): + break backlog_len = self.client.quotes.qsize() - if message is not None and len(message) > 0 and len(message) >= message[0] * 24: #sanity check on length. Should be at least as long as the smallest message times the number of messages it says it has. + if not isinstance(message, (bytes, bytearray, memoryview)): + self.client.logger.error("Ignored non-binary websocket message") + continue + try: + if len(message) == 0 or len(message) < 1 + message[0] * 24: + raise ValueError("Equities message is shorter than its item count") items_in_message = message[0] start_index = 1 for i in range(0, items_in_message): - start_index = self.parse_message(message, start_index, backlog_len) \ No newline at end of file + if self._stop_event.is_set(): + break + if start_index + 2 > len(message): + raise ValueError("Equities message item header is truncated") + message_length = message[start_index + 1] + if message_length < 24 or start_index + message_length > len(message): + raise ValueError("Equities message item length is invalid") + start_index = self.parse_message(message, start_index, backlog_len) + except Exception as error: + self.client.logger.error(f"Ignored malformed equities message: {repr(error)}") diff --git a/intriniorealtime/options_client.py b/intriniorealtime/options_client.py index 3815b70..073ac6b 100644 --- a/intriniorealtime/options_client.py +++ b/intriniorealtime/options_client.py @@ -9,7 +9,11 @@ from collections.abc import Callable from enum import IntEnum, unique +from ._websocket import close_websocket_app, should_abort_handshake, _redact_api_key + _SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] +_CONNECT_TIMEOUT_SECONDS = 30 +_STABLE_SESSION_SECONDS = 60.0 _EMPTY_STRING = "" _OPTIONS_TRADE_MESSAGE_SIZE = 72 # 61 used + 11 pad _OPTIONS_QUOTE_MESSAGE_SIZE = 52 # 48 used + 4 pad @@ -18,10 +22,6 @@ _NAN = float("NAN") _stopFlag: threading.Event = threading.Event() -_dataMsgLock: threading.Lock = threading.Lock() -_dataMsgCount: int = 0 -_txtMsgLock: threading.Lock = threading.Lock() -_txtMsgCount: int = 0 _logHandler: logging.Logger = logging.StreamHandler() _logHandler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) _log: logging.Logger = logging.getLogger('intrinio_realtime_options') @@ -31,15 +31,21 @@ def log(message: str): _log.info(message) -def do_backoff(fn: Callable[[None], bool]): +def do_backoff(fn: Callable[[], bool], stop_flag: threading.Event = None): + flag: threading.Event = stop_flag if stop_flag is not None else _stopFlag i: int = 0 - backoff: int = _SELF_HEAL_BACKOFFS[i] - success: bool = fn() - while (not success): - time.sleep(backoff) + while not flag.is_set(): + try: + if fn(): + return + except Exception as e: + _log.warning("Websocket - Attempt failed: {0}".format(e)) + if flag.is_set(): + return + backoff: int = _SELF_HEAL_BACKOFFS[i] i = min(i + 1, len(_SELF_HEAL_BACKOFFS) - 1) - backoff = _SELF_HEAL_BACKOFFS[i] - success = fn() + if flag.wait(timeout=backoff): + return @unique class Providers(IntEnum): @@ -270,27 +276,36 @@ def _get_option_mask(use_on_trade: bool, use_on_quote: bool, use_on_refresh: boo mask |= 0b1000 return mask -class _WebSocket(websocket.WebSocketApp): +class _WebSocket: + on_cont_message = None + def __init__(self, ws_url: str, ws_lock: threading.Lock, + channels_lock: threading.RLock, + sent_channels: set[str], + stats: dict[str, int], + stats_lock: threading.Lock, worker_threads: list[threading.Thread], - get_channels: Callable[[None], set[tuple[str, bool]]], - get_token: Callable[[None], str], + get_channels: Callable[[], set], + get_token: Callable[[], str], get_url: Callable[[str], str], use_on_trade: bool, use_on_quote: bool, use_on_refresh: bool, use_on_ua: bool, - data_queue: queue.Queue): - super().__init__(ws_url, on_open=self.__on_open, on_close=self.__on_close, on_data=self.__on_data, on_error=self.__on_error) + data_queue: queue.Queue, + stop_flag: threading.Event, + worker_factory: Callable[[int], threading.Thread] = None): + self.url: str = ws_url self.__wsLock: threading.Lock = ws_lock - self.__continuation_queue = queue.Queue(100) - self.__continuation_lock: threading.Lock = threading.Lock() - self.__currently_continuing: bool = False + self.__channels_lock: threading.RLock = channels_lock + self.__sent_channels: set[str] = sent_channels + self.__stats: dict[str, int] = stats + self.__stats_lock: threading.Lock = stats_lock self.__worker_threads: list[threading.Thread] = worker_threads - self.__get_channels: Callable[[None], set[tuple[str, bool]]] = get_channels - self.__get_token: Callable[[None], str] = get_token + self.__get_channels: Callable[[], set] = get_channels + self.__get_token: Callable[[], str] = get_token self.__get_url: Callable[[str], str] = get_url self.__use_on_trade: bool = use_on_trade self.__use_on_quote: bool = use_on_quote @@ -300,105 +315,252 @@ def __init__(self, self.__is_reconnecting: bool = False self.__last_reset: float = time.time() self.isReady: bool = False + self.__handshake_event: threading.Event = threading.Event() + self.__attempt_cancelled: threading.Event = threading.Event() + self.__connect_deadline: float = float("inf") + self.__backoff_index: int = 0 + self.__connect_generation: int = 0 + self.__connect_lock: threading.Lock = threading.Lock() + self.__app_lock: threading.Lock = threading.Lock() + self.__app = None + self.__session_opened: bool = False + self.__session_opened_at: float = 0.0 + self.__stop_flag: threading.Event = stop_flag + self.__worker_factory = worker_factory + + def __close_leftover_sock(self): + with self.__app_lock: + app = self.__app + close_websocket_app(app) def __on_open(self, ws): + rejected = False + with self.__connect_lock: + if (self.__attempt_cancelled.is_set() + or self.__stop_flag.is_set() + or time.monotonic() >= self.__connect_deadline): + self.__attempt_cancelled.set() + rejected = True + else: + self.__session_opened = True + self.__session_opened_at = time.monotonic() + self.__handshake_event.set() + if rejected: + close_websocket_app(ws) + return _log.info("Websocket - Connected") - self.__wsLock.acquire() - try: - self.isReady = True - self.__is_reconnecting = False - for worker in self.__worker_threads: - if not worker.is_alive(): - worker.start() - finally: - self.__wsLock.release() - if self.__get_channels and callable(self.__get_channels): - channels: set[str] = self.__get_channels() - if channels and (len(channels) > 0): - for symbol in channels: - symbol_bytes = bytes(symbol, 'utf-8') - message: bytes = bytearray(len(symbol_bytes) + 2) - message[0] = 74 # join code - message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_ua) - message[2:] = symbol_bytes - if self.isReady: - _log.info("Websocket - Joining channel: {0}".format(symbol)) - self.send_binary(message) - - def __try_reconnect(self) -> bool: - _log.info("Websocket - Reconnecting...") - if self.isReady: - return True - else: - with self.__wsLock: - self.__is_reconnecting = True - token: str = self.__get_token(None) - super().url = self.__get_url(token) - self.start() - return False + rejected = False + with self.__channels_lock: + if self.__attempt_cancelled.is_set() or self.__stop_flag.is_set(): + self.__attempt_cancelled.set() + rejected = True + else: + self.__sent_channels.clear() + self.__wsLock.acquire() + try: + for index, worker in enumerate(self.__worker_threads): + if worker.is_alive(): + continue + if worker.ident is not None: + if self.__worker_factory is None: + raise RuntimeError("Stopped worker thread cannot be restarted") + worker = self.__worker_factory(index) + self.__worker_threads[index] = worker + worker.start() + self.isReady = True + self.__is_reconnecting = False + except Exception: + self.isReady = False + self.__sent_channels.clear() + close_websocket_app(ws) + raise + finally: + self.__wsLock.release() + if self.__get_channels and callable(self.__get_channels): + channels: set[str] = self.__get_channels() + if channels and (len(channels) > 0): + for symbol in channels: + symbol_bytes = bytes(symbol, 'utf-8') + message: bytes = bytearray(len(symbol_bytes) + 2) + message[0] = 74 # join code + message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_ua) + message[2:] = symbol_bytes + if self.isReady: + _log.info("Websocket - Joining channel: {0}".format(symbol)) + try: + self.send_binary(message) + self.__sent_channels.add(symbol) + except Exception: + self.__wsLock.acquire() + try: + self.isReady = False + finally: + self.__wsLock.release() + self.__sent_channels.clear() + close_websocket_app(ws) + raise + if rejected: + close_websocket_app(ws) def __on_close(self, ws, closeStatusCode, closeMsg): + _log.info("Websocket - Closed - {0}: {1}".format(closeStatusCode, closeMsg)) + with self.__channels_lock: + self.__sent_channels.clear() self.__wsLock.acquire() try: - if (not self.__is_reconnecting): - _log.info("Websocket - Closed - {0}: {1}".format(closeStatusCode, closeMsg)) - self.isReady = False - if (not _stopFlag.is_set()): - do_backoff(self.__try_reconnect) + self.isReady = False finally: self.__wsLock.release() def __on_error(self, ws, error): _log.error("Websocket - Error - {0}".format(error)) - def __stitch(self): - full = None - while not self.__continuation_queue.empty(): - partial = self.__continuation_queue.get(True, 1) - if full is None: - full = partial - else: - full = full.join(partial) - return full - - def __on_data(self, ws, data, code, is_last): #continueFlag - If 0, the data continues + def __on_data(self, ws, data, code, is_last): if code == websocket.ABNF.OPCODE_BINARY: - with _dataMsgLock: - global _dataMsgCount - _dataMsgCount += 1 - - if self.__currently_continuing or is_last == 0: # we're either in the middle of a continue, or we're starting a continue - with self.__continuation_lock: - if self.__currently_continuing or is_last == 0: #check lock check - #_log.info(f"Received partial message (hex): {data.hex()}") - self.__continuation_queue.put(data) - if is_last != 0: - full_message = self.__stitch() - self.__data_queue.put(full_message) - self.__currently_continuing = True if is_last == 0 else False # we're in the middle of a continue, but this is the last message, so remove flag - else: # We're not in the middle of a continue, and this isn't marked as multi-part, so this is a full message by itself. - self.__data_queue.put(data) - else: # We're not in the middle of a continue, and this isn't marked as multi-part, so this is a full message by itself. - self.__data_queue.put(data) + with self.__stats_lock: + self.__stats["data"] += 1 + self.__data_queue.put(data) else: _log.debug("Websocket - Message received") - with _txtMsgLock: - global _txtMsgCount - _txtMsgCount += 1 + with self.__stats_lock: + self.__stats["text"] += 1 _log.error("Error received: {0}".format(data)) def start(self): - super().run_forever(skip_utf8_validation=True) + while not self.__stop_flag.is_set(): + handshake_event = threading.Event() + self.__handshake_event = handshake_event + attempt_cancelled = threading.Event() + self.__attempt_cancelled = attempt_cancelled + connect_deadline = time.monotonic() + _CONNECT_TIMEOUT_SECONDS + self.__connect_deadline = connect_deadline + with self.__connect_lock: + self.__connect_generation += 1 + generation = self.__connect_generation + self.__session_opened = False + self.__wsLock.acquire() + try: + self.isReady = False + finally: + self.__wsLock.release() + self.__close_leftover_sock() + + # Leave on_cont_message unset so websocket-client aggregates fragmented messages. + app = websocket.WebSocketApp( + self.url, + on_open=self.__on_open, + on_close=self.__on_close, + on_data=self.__on_data, + on_error=self.__on_error, + ) + with self.__app_lock: + if self.__stop_flag.is_set(): + close_websocket_app(app) + return + self.__app = app + + def watchdog( + this_generation=generation, + handshake_event=handshake_event, + cancelled=attempt_cancelled, + deadline=connect_deadline, + attempt_app=app, + ): + if should_abort_handshake( + handshake_event, + self.__stop_flag, + this_generation, + lambda: self.__connect_generation, + deadline, + ): + with self.__connect_lock: + if this_generation != self.__connect_generation: + return + _log.warning("Websocket - Connect timed out after {0}s".format(_CONNECT_TIMEOUT_SECONDS)) + cancelled.set() + close_websocket_app(attempt_app) + + threading.Thread(target=watchdog, daemon=True).start() + if self.__stop_flag.is_set(): + close_websocket_app(app) + return + try: + app.run_forever(skip_utf8_validation=True) + except Exception as e: + _log.warning("Websocket - Attempt failed: {0}".format(e)) + finally: + # The attempt is over even if the websocket never opened. Wake + # its watchdog so it cannot fire later during reconnect backoff. + handshake_event.set() + with self.__app_lock: + if self.__app is app: + self.__app = None + + self.__wsLock.acquire() + try: + self.isReady = False + finally: + self.__wsLock.release() + + if self.__stop_flag.is_set(): + return + + _log.info("Websocket - Reconnecting...") + skip_wait: bool = ( + self.__session_opened + and (time.monotonic() - self.__session_opened_at) >= _STABLE_SESSION_SECONDS + ) + if skip_wait: + self.__backoff_index = 0 + token_ready: bool = False + while not self.__stop_flag.is_set(): + if skip_wait: + skip_wait = False + else: + backoff: int = _SELF_HEAL_BACKOFFS[self.__backoff_index] + self.__backoff_index = min(self.__backoff_index + 1, len(_SELF_HEAL_BACKOFFS) - 1) + if self.__stop_flag.wait(timeout=backoff): + return + try: + token: str = self.__get_token() + if token: + self.url = self.__get_url(token) + token_ready = True + break + _log.warning("Websocket - Reconnect attempt failed: no token") + except Exception as e: + _log.warning("Websocket - Reconnect attempt failed: {0}".format(e)) + if not token_ready: + return # super().run_forever(ping_interval = 5, ping_timeout = 2, skip_utf8_validation = True) def stop(self): - super().close() + self.__close_leftover_sock() + + def close(self): + self.stop() + + # SDK-owned callback entry points; these do not expose or mutate + # websocket-client state. + def on_open(self, ws): + self.__on_open(ws) + + def on_data(self, ws, data, code, is_last): + self.__on_data(ws, data, code, is_last) def send(self, message: str): - super().send(message, websocket.ABNF.OPCODE_TEXT) + with self.__app_lock: + app = self.__app + if app is None: + raise RuntimeError("Websocket is not connected") + app.send(message, websocket.ABNF.OPCODE_TEXT) def send_binary(self, message: bytes): - super().send(message, websocket.ABNF.OPCODE_BINARY) + with self.__app_lock: + app = self.__app + if app is None: + raise RuntimeError("Websocket is not connected") + app.send(message, websocket.ABNF.OPCODE_BINARY) def reset(self): self.__last_reset = time.time() @@ -406,11 +568,13 @@ def reset(self): class Config: def __init__(self, api_key: str, provider: Providers, num_threads: int = 4, log_level: LogLevel = LogLevel.INFO, manual_ip_address: str = None, symbols: set[str] = None, delayed: bool = False): + if isinstance(num_threads, bool) or not isinstance(num_threads, int) or num_threads < 1: + raise ValueError("num_threads must be a positive integer") self.api_key: str = api_key self.provider: Providers = provider self.num_threads: int = num_threads self.manual_ip_address: str = manual_ip_address - self.symbols: list[str] = symbols + self.symbols: set[str] = symbols self.log_level: LogLevel = log_level self.delayed: bool = delayed @@ -500,19 +664,23 @@ def _thread_fn(index: int, data: queue.Queue, on_trade: Callable[[OptionsTrade], None], on_quote: Callable[[OptionsQuote], None] = None, on_refresh: Callable[[OptionsRefresh], None] = None, - on_unusual_activity: Callable[[OptionsUnusualActivity], None] = None): + on_unusual_activity: Callable[[OptionsUnusualActivity], None] = None, + stop_flag: threading.Event = None): + flag: threading.Event = stop_flag if stop_flag is not None else _stopFlag _log.debug("Starting worker thread {0}".format(index)) datum: bytes = None count: int = 0 start_index: int = 1 msg_type: int = 0 message: bytes = None - while not _stopFlag.is_set(): + while not flag.is_set(): try: datum = data.get(True, 1.0) count = datum[0] start_index = 1 for _ in range(count): + if flag.is_set(): + break msg_type = datum[start_index + 22] if msg_type == 1: # Quote message: bytes = datum[start_index:(start_index + _OPTIONS_QUOTE_MESSAGE_SIZE)] @@ -621,7 +789,9 @@ def _thread_fn(index: int, data: queue.Queue, continue except Exception as e: _log.error(f"Worker thread {index} Exception {e}") - _log.error(f"\tCurrent count: {count}\r\n\tCurrent start_index: {start_index}\r\n\tCurrent msg_type: {msg_type}\r\n\tFull message (hex): {datum.hex()}\r\n\tScoped message: {message.hex()}") + datum_hex = datum.hex() if isinstance(datum, (bytes, bytearray)) else repr(datum) + message_hex = message.hex() if isinstance(message, (bytes, bytearray)) else repr(message) + _log.error(f"\tCurrent count: {count}\r\n\tCurrent start_index: {start_index}\r\n\tCurrent msg_type: {msg_type}\r\n\tFull message (hex): {datum_hex}\r\n\tScoped message: {message_hex}") continue _log.debug("Worker thread {0} stopped".format(index)) @@ -675,23 +845,57 @@ def __init__(self, config: Config, on_trade: Callable[[OptionsTrade], None], on_ self.__manualIP: str = config.manual_ip_address self.__token: tuple[str, float] = (None, 0.0) self.__webSocket: _WebSocket = None - if config.symbols and (isinstance(config.symbols, list)) and (len(config.symbols) > 0): + if config.symbols and isinstance(config.symbols, (list, set)): self.__channels: set[str] = set((_transform_contract_to_new(symbol)) for symbol in config.symbols) else: self.__channels: set[str] = set() self.__data: queue.Queue = queue.Queue() self.__t_lock: threading.Lock = threading.Lock() self.__ws_lock: threading.Lock = threading.Lock() - self.__worker_threads: list[threading.Thread] = [threading.Thread( + self.__channels_lock: threading.RLock = threading.RLock() + self.__sent_channels: set[str] = set() + self.__stats: dict[str, int] = {"data": 0, "text": 0} + self.__stats_lock: threading.Lock = threading.Lock() + self.__lifecycle_lock: threading.RLock = threading.RLock() + self.__on_trade = on_trade + self.__on_quote = on_quote + self.__on_refresh = on_refresh + self.__on_unusual_activity = on_unusual_activity + if (isinstance(config.num_threads, bool) + or not isinstance(config.num_threads, int) + or config.num_threads < 1): + raise ValueError("num_threads must be a positive integer") + self.__num_threads: int = config.num_threads + self._stop_event: threading.Event = threading.Event() + self._stop_generation: int = 0 + self.__worker_threads: list[threading.Thread] = self.__make_worker_threads() + self.__socket_thread: threading.Thread = None + self.__is_started: bool = False + self.__is_starting: bool = False + self.__is_stopping: bool = False + self.__start_complete: threading.Event = threading.Event() + self.__start_complete.set() + self.__start_result = {"error": None} + self.__stopping_thread = None + self.__stop_complete: threading.Event = threading.Event() + self.__stop_complete.set() + self.__stop_result = {"error": None} + _log.setLevel(config.log_level) + + def __make_worker_threads(self) -> list[threading.Thread]: + return [self.__make_worker_thread(i, self.__data, self._stop_event) + for i in range(self.__num_threads)] + + def __make_worker_thread(self, index: int, data_queue: queue.Queue, + stop_event: threading.Event) -> threading.Thread: + return threading.Thread( group=None, target=_thread_fn, - args=(i, self.__data, on_trade, on_quote, on_refresh, on_unusual_activity), + args=(index, data_queue, self.__on_trade, self.__on_quote, + self.__on_refresh, self.__on_unusual_activity, stop_event), kwargs={}, daemon=True - ) for i in range(config.num_threads)] - self.__socket_thread: threading.Thread = None - self.__is_started: bool = False - _log.setLevel(config.log_level) + ) def __all_ready(self) -> bool: self.__ws_lock.acquire() @@ -743,67 +947,83 @@ def __try_set_token(self) -> bool: _log.error("Authorization Failure: The request timed out.") return False except requests.exceptions.ConnectionError as err: - _log.error("Authorization Failure: {0}".format(err)) + _log.error("Authorization Failure: %s", _redact_api_key(err, self.__apiKey)) return False def __get_token(self) -> str: self.__t_lock.acquire() try: if ((time.time() - self.__token[1]) > (60 * 60 * 24)): # 60sec/min * 60min/hr * 24hrs = 1 day - do_backoff(self.__try_set_token) + do_backoff(self.__try_set_token, self._stop_event) return self.__token[0] finally: self.__t_lock.release() + def __refresh_token(self) -> str: + self.__t_lock.acquire() + try: + self.__token = (None, 0.0) + finally: + self.__t_lock.release() + return self.__get_token() + def __get_channels(self) -> set[str]: - return self.__channels + with self.__channels_lock: + return self.__channels.copy() def __join(self, symbol: str): transformed_symbol: str = _transform_contract_to_new(symbol) - if transformed_symbol not in self.__channels: - self.__channels.add(transformed_symbol) - symbol_bytes = bytes(transformed_symbol, 'utf-8') - message: bytes = bytearray(len(symbol_bytes)+2) - message[0] = 74 # join code - message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity) - message[2:] = symbol_bytes - if self.__webSocket.isReady: - _log.info("Websocket - Joining channel: {0}".format(transformed_symbol)) - self.__webSocket.send_binary(message) - - def __leave(self, symbol: str): + with self.__channels_lock: + if transformed_symbol not in self.__channels: + self.__channels.add(transformed_symbol) + if transformed_symbol not in self.__sent_channels: + symbol_bytes = bytes(transformed_symbol, 'utf-8') + message: bytes = bytearray(len(symbol_bytes)+2) + message[0] = 74 # join code + message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity) + message[2:] = symbol_bytes + ws = self.__webSocket + if ws is not None and ws.isReady: + _log.info("Websocket - Joining channel: {0}".format(transformed_symbol)) + ws.send_binary(message) + self.__sent_channels.add(transformed_symbol) + + def __leave(self, symbol: str, forget: bool = True): transformed_symbol: str = _transform_contract_to_new(symbol) - if transformed_symbol in self.__channels: - self.__channels.remove(transformed_symbol) + with self.__channels_lock: + if transformed_symbol not in self.__channels: + return symbol_bytes = bytes(transformed_symbol, 'utf-8') message: bytes = bytearray(len(symbol_bytes) + 2) message[0] = 76 # leave code message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity) message[2:] = symbol_bytes - if self.__webSocket.isReady: - _log.info("Websocket - Leaving channel: {0}".format(transformed_symbol)) - self.__webSocket.send_binary(message) + try: + ws = self.__webSocket + if ws is not None and ws.isReady: + _log.info("Websocket - Leaving channel: {0}".format(transformed_symbol)) + ws.send_binary(message) + finally: + self.__sent_channels.discard(transformed_symbol) + if forget: + self.__channels.discard(transformed_symbol) def join(self, *symbols): - if self.__is_started: - while not self.__all_ready(): - time.sleep(1.0) for (symbol) in symbols: self.__join(symbol) def join_firehose(self): - if "$FIREHOSE" in self.__channels: - _log.warn("This client has already joined the firehose channel") - else: - if self.__is_started: - while not self.__all_ready(): - time.sleep(1.0) - self.__join("$FIREHOSE") + with self.__channels_lock: + if "$FIREHOSE" in self.__channels: + _log.warn("This client has already joined the firehose channel") + return + self.__join("$FIREHOSE") def leave(self, *symbols): if not symbols: _log.info("Leaving all channels") - channels: set[str] = self.__channels.copy() + with self.__channels_lock: + channels: set[str] = self.__channels.copy() for (symbol) in channels: self.__leave(symbol) symbol_set: set[str] = set(symbols) @@ -811,61 +1031,277 @@ def leave(self, *symbols): self.__leave(sym) def leave_firehose(self): - if "$FIREHOSE" in self.__channels: + with self.__channels_lock: + joined = "$FIREHOSE" in self.__channels + if joined: self.__leave("$FIREHOSE") - def __socket_start_fn(self, token: str): + def __socket_start_fn(self, token: str, stop_event: threading.Event, + worker_threads: list[threading.Thread], data_queue: queue.Queue): + if not token: + _log.error("Websocket - Missing token") + return _log.info("Websocket - Connecting...") ws_url: str = self.__get_web_socket_url(token) - self.__webSocket = _WebSocket(ws_url, - self.__ws_lock, - self.__worker_threads, - self.__get_channels, - self.__get_token, - self.__get_web_socket_url, - self.__use_on_trade, - self.__use_on_quote, - self.__use_on_refresh, - self.__use_on_unusual_activity, - self.__data) - self.__webSocket.start() + if stop_event.is_set(): + return + web_socket = _WebSocket(ws_url, + self.__ws_lock, + self.__channels_lock, + self.__sent_channels, + self.__stats, + self.__stats_lock, + worker_threads, + self.__get_channels, + self.__refresh_token, + self.__get_web_socket_url, + self.__use_on_trade, + self.__use_on_quote, + self.__use_on_refresh, + self.__use_on_unusual_activity, + data_queue, + stop_event, + lambda index: self.__make_worker_thread( + index, data_queue, stop_event)) + self.__webSocket = web_socket + web_socket.start() def start(self): if (not (self.__use_on_trade or self.__use_on_quote or self.__use_on_refresh or self.__use_on_unusual_activity)): raise ValueError("You must set at least one callback method before starting client") - token: str = self.__get_token() - self.__ws_lock.acquire() + with self.__lifecycle_lock: + if (self.__is_started + and self.__socket_thread is not None + and self.__socket_thread.is_alive() + and not self._stop_event.is_set()): + return + if self.__is_stopping: + raise RuntimeError("Cannot start while client is stopping") + if self.__is_starting: + start_complete = self.__start_complete + start_result = self.__start_result + wait_for_start = True + else: + wait_for_start = False + self.__is_starting = True + start_complete = threading.Event() + start_result = {"error": None} + self.__start_complete = start_complete + self.__start_result = start_result + stop_generation = self._stop_generation + if wait_for_start: + start_complete.wait() + if start_result["error"] is not None: + raise RuntimeError( + "Concurrent start failed: {0}".format(start_result["error"]) + ) from start_result["error"] + return + + failure = None try: - self.__socket_thread = threading.Thread( - group=None, - target=self.__socket_start_fn, - args=(token,), - kwargs={}, - daemon=True - ) + while True: + leftover = self.__socket_thread + if leftover is not None and leftover.is_alive(): + close_websocket_app(self.__webSocket) + leftover.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) + if leftover.is_alive(): + with self.__lifecycle_lock: + if self._stop_generation != stop_generation: + raise RuntimeError("Cannot start while client is stopping") + raise RuntimeError("Previous socket thread did not exit") + if any(worker.is_alive() for worker in self.__worker_threads): + raise RuntimeError("Previous worker threads did not exit") + with self.__lifecycle_lock: + if self._stop_generation != stop_generation: + raise RuntimeError("Cannot start while client is stopping") + self._stop_event = threading.Event() + stop_event = self._stop_event + self.__data = queue.Queue() + data_queue = self.__data + self.__worker_threads = self.__make_worker_threads() + worker_threads = self.__worker_threads + token: str = self.__get_token() + if stop_event.is_set(): + raise RuntimeError("Cannot start while client is stopping") + if not token: + if stop_event.wait(timeout=_SELF_HEAL_BACKOFFS[0]): + raise RuntimeError("Cannot start while client is stopping") + continue + with self.__lifecycle_lock: + if (self._stop_generation != stop_generation + or self._stop_event is not stop_event + or stop_event.is_set()): + raise RuntimeError("Cannot start while client is stopping") + socket_thread = threading.Thread( + group=None, + target=self.__socket_start_fn, + args=(token, stop_event, worker_threads, data_queue), + kwargs={}, + daemon=True + ) + self.__socket_thread = socket_thread + self.__is_started = True + try: + socket_thread.start() + except Exception: + self.__is_started = False + raise + if (self._stop_generation != stop_generation + or self._stop_event is not stop_event + or stop_event.is_set()): + self.__is_started = False + raise RuntimeError("Cannot start while client is stopping") + return + except BaseException as error: + failure = error + raise finally: - self.__ws_lock.release() - self.__socket_thread.start() - self.__is_started = True + with self.__lifecycle_lock: + self.__is_starting = False + start_result["error"] = failure + start_complete.set() + + def __publish_stop_result(self, stop_complete: threading.Event, + stop_result: dict, surviving_threads: list, + failure: BaseException = None): + with self.__lifecycle_lock: + self.__is_started = bool(surviving_threads) + self.__is_stopping = False + self.__stopping_thread = None + stop_result["error"] = str(failure) if failure is not None else None + stop_complete.set() + if failure is not None: + _log.error("Stop incomplete: {0}".format(failure)) + else: + _log.info("Stopped") + + def __finish_stop_after_owned_thread_exits( + self, owned_thread: threading.Thread, + candidate_threads: list[threading.Thread], + stop_complete: threading.Event, stop_result: dict, + failure: BaseException = None): + try: + owned_thread.join() + except BaseException as error: + if failure is None: + failure = error + surviving_threads = [ + thread for thread in candidate_threads if thread.is_alive() + ] + if failure is None and surviving_threads: + failure = RuntimeError("Client stop incomplete: threads did not exit") + self.__publish_stop_result( + stop_complete, stop_result, surviving_threads, failure) def stop(self): - _log.info("Stopping...") - if len(self.__channels) > 0: - self.leave() - time.sleep(1.0) - self.__ws_lock.acquire() + current_thread = threading.current_thread() + with self.__lifecycle_lock: + if self.__is_stopping: + owned_threads = self.__worker_threads + ( + [self.__socket_thread] + if self.__socket_thread is not None else [] + ) + if (self.__stopping_thread is current_thread + or current_thread in owned_threads): + # An owned thread must return to its run loop before the + # active shutdown can join it. Waiting here would make the + # shutdown wait on its own completion event. + return + stop_complete = self.__stop_complete + stop_result = self.__stop_result + wait_for_stop = True + else: + wait_for_stop = False + self.__is_stopping = True + self.__stopping_thread = current_thread + stop_complete = threading.Event() + stop_result = {"error": None} + self.__stop_complete = stop_complete + self.__stop_result = stop_result + _log.info("Stopping...") + self._stop_generation += 1 + self._stop_event.set() + ws = self.__webSocket + worker_threads = self.__worker_threads + socket_thread = self.__socket_thread + shutdown_deadline = time.monotonic() + _CONNECT_TIMEOUT_SECONDS + 1 + if wait_for_stop: + stop_complete.wait() + if stop_result["error"] is not None: + raise RuntimeError(stop_result["error"]) + return + + candidate_threads = [ + worker for worker in worker_threads if worker.ident is not None + ] + if socket_thread is not None and socket_thread.ident is not None: + candidate_threads.append(socket_thread) + surviving_threads = [] + failure = None + deferred_completion = False try: - self.__webSocket.isReady = False + try: + if ws is not None: + self.__ws_lock.acquire() + try: + ws.isReady = False + finally: + self.__ws_lock.release() + ws.stop() + with self.__channels_lock: + self.__sent_channels.clear() + finally: + self.__ws_lock.acquire() + try: + if ws is not None: + ws.isReady = False + finally: + self.__ws_lock.release() + if ws is not None: + ws.stop() + for i in range(len(worker_threads)): + worker = worker_threads[i] + if worker.ident is not None and worker is not current_thread: + worker.join(timeout=max(0.0, shutdown_deadline - time.monotonic())) + if (socket_thread is not None + and socket_thread.ident is not None + and socket_thread is not current_thread): + socket_thread.join(timeout=max(0.0, shutdown_deadline - time.monotonic())) + + surviving_threads = [thread for thread in candidate_threads if thread.is_alive()] + except BaseException as error: + failure = error finally: - self.__ws_lock.release() - _stopFlag.set() - self.__webSocket.stop() - for i in range(len(self.__worker_threads)): - self.__worker_threads[i].join() - _log.debug("Worker thread {0} joined".format(i)) - self.__socket_thread.join() - _log.debug("Socket thread joined") - _log.info("Stopped") + # Recheck after all joins and cleanup so threads that exited while a + # later thread was being joined are not retained as false survivors. + surviving_threads = [thread for thread in candidate_threads if thread.is_alive()] + owned_current_thread = ( + current_thread if current_thread in surviving_threads else None + ) + other_survivors = [ + thread for thread in surviving_threads + if thread is not owned_current_thread + ] + if failure is None and other_survivors: + failure = RuntimeError("Client stop incomplete: threads did not exit") + if owned_current_thread is not None: + deferred_completion = True + threading.Thread( + target=self.__finish_stop_after_owned_thread_exits, + args=(owned_current_thread, candidate_threads, + stop_complete, stop_result, failure), + daemon=True, + ).start() + else: + self.__publish_stop_result( + stop_complete, stop_result, surviving_threads, failure) + if failure is not None: + raise failure + if deferred_completion: + _log.debug("Stop completion deferred until callback thread exits") def get_stats(self) -> tuple[int, int, int]: - return _dataMsgCount, _txtMsgCount, self.__data.qsize() + with self.__stats_lock: + data_count = self.__stats["data"] + text_count = self.__stats["text"] + return data_count, text_count, self.__data.qsize() diff --git a/setup.py b/setup.py index 5b6d081..8a7d21d 100644 --- a/setup.py +++ b/setup.py @@ -7,7 +7,7 @@ def readme(): setup( name = 'intriniorealtime', packages = ['intriniorealtime'], - version = '6.3.0', + version = '6.4.0', author = 'Intrinio Python SDK for Real-Time Stock Prices', author_email = 'success@intrinio.com', url = 'https://intrinio.com', @@ -16,7 +16,7 @@ def readme(): long_description_content_type = 'text/markdown', install_requires = ['requests>=2.26.0','websocket-client>=1.2.1','wsaccel>=0.6.3', 'intrinio-sdk>=6.26.0'], python_requires = '~=3.10', - download_url = 'https://github.com/intrinio/intrinio-realtime-python-sdk/archive/v6.3.0.tar.gz', + download_url = 'https://github.com/intrinio/intrinio-realtime-python-sdk/archive/v6.4.0.tar.gz', keywords = ['realtime','stock prices','intrinio','stock market','stock data','financial'], classifiers = [ 'Intended Audience :: Financial and Insurance Industry', diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_equities_client.py b/tests/test_equities_client.py new file mode 100644 index 0000000..189315e --- /dev/null +++ b/tests/test_equities_client.py @@ -0,0 +1,915 @@ +import logging +import queue +import threading +import time +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import websocket + +from intriniorealtime.equities_client import ( + CONNECT_TIMEOUT_SECONDS, + STABLE_SESSION_SECONDS, + EquitiesQuoteHandler, + EquitiesQuoteReceiver, + IntrinioRealtimeEquitiesClient, +) + + +def _make_client(): + with patch.object(EquitiesQuoteHandler, "start"): + return IntrinioRealtimeEquitiesClient( + {"api_key": "test-key", "provider": "IEX"}, + lambda *args: None, + lambda *args: None, + ) + + +class EquitiesReconnectTests(unittest.TestCase): + def setUp(self): + logging.disable(logging.CRITICAL) + + def tearDown(self): + logging.disable(logging.NOTSET) + + def test_default_stable_session_requires_sustained_health(self): + self.assertGreaterEqual(STABLE_SESSION_SECONDS, 60) + + def test_on_connect_rejoins_channels_after_reconnect(self): + client = _make_client() + client.channels = {"AAPL"} + client.joined_channels = {"AAPL"} + client.ready = False + client.ws = Mock() + + client.on_connect() + + self.assertTrue(client.ready) + self.assertEqual(client.joined_channels, {"AAPL"}) + client.ws.send.assert_called() + + def test_on_connect_does_not_skip_join_when_joined_channels_are_stale(self): + client = _make_client() + client.channels = {"AAPL", "MSFT"} + client.joined_channels = {"AAPL", "MSFT"} + client.ready = True + client.ws = Mock() + + client.refresh_channels() + client.ws.send.assert_not_called() + + client.on_connect() + self.assertEqual(client.ws.send.call_count, 2) + + def test_on_connect_does_not_restore_ready_after_disconnect(self): + client = _make_client() + client.ws = Mock() + client._stop_event.set() + + client.on_connect() + + self.assertFalse(client.ready) + client.ws.send.assert_not_called() + + def test_watchdog_aborts_app_when_handshake_times_out(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + aborted = [] + + def fake_run_forever(app, *args, **kwargs): + time.sleep(0.12) + receiver.enabled = False + client._stop_event.set() + + with patch("intriniorealtime.equities_client.close_websocket_app", side_effect=lambda app: aborted.append(app)), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("websocket.setdefaulttimeout"), \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever): + receiver.run() + + self.assertTrue(aborted) + + def test_run_does_not_change_websocket_client_global_timeout(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + + def fake_run_forever(app, *args, **kwargs): + app.on_open(app) + receiver.enabled = False + client._stop_event.set() + + with patch("websocket.setdefaulttimeout") as set_timeout, \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever): + receiver.run() + + set_timeout.assert_not_called() + + def test_late_open_after_disconnect_is_aborted_before_becoming_ready(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + aborted = [] + + def fake_run_forever(app, *args, **kwargs): + receiver.enabled = False + client._stop_event.set() + app.on_open(app) + + with patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.close_websocket_app", side_effect=aborted.append): + receiver.run() + + self.assertFalse(client.ready) + self.assertTrue(aborted) + + def test_late_open_after_timeout_is_aborted_before_becoming_ready(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + aborted = [] + + def fake_run_forever(app, *args, **kwargs): + time.sleep(0.08) + app.on_open(app) + receiver.enabled = False + + with patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.03), \ + patch("intriniorealtime.equities_client.should_abort_handshake", return_value=False), \ + patch("intriniorealtime.equities_client.close_websocket_app", side_effect=aborted.append): + receiver.run() + + self.assertFalse(client.ready) + self.assertTrue(aborted) + + def test_refresh_token_failure_backs_off_and_skips_socket_until_auth_succeeds(self): + client = _make_client() + client.token = "stale" + client.do_backoff = Mock() + attempts = {"n": 0} + + def refresh(): + attempts["n"] += 1 + if attempts["n"] < 3: + raise RuntimeError("Auth failed") + + client.refresh_token = Mock(side_effect=refresh) + receiver = EquitiesQuoteReceiver(client) + run_forever_calls = [] + + def fake_run_forever(app, *args, **kwargs): + run_forever_calls.append(1) + if len(run_forever_calls) >= 2: + receiver.enabled = False + client._stop_event.set() + + with patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.close_websocket_app"): + receiver.run() + + self.assertEqual(len(run_forever_calls), 2) + self.assertEqual(client.refresh_token.call_count, 3) + self.assertGreaterEqual(client.do_backoff.call_count, 3) + + def test_disconnect_then_connect_starts_a_new_receiver(self): + client = _make_client() + client.refresh_token = Mock() + old = Mock() + old.is_alive.return_value = True + old.enabled = True + + def join_and_exit(timeout=None): + old.is_alive.return_value = False + + old.join.side_effect = join_and_exit + client.quote_receiver = old + client.ws = Mock() + + with patch("intriniorealtime.equities_client.close_websocket_app"): + client.disconnect() + + self.assertFalse(old.enabled) + old.join.assert_called() + self.assertIsNone(client.quote_receiver) + self.assertTrue(client._stop_event.is_set()) + + new_receiver = Mock() + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", return_value=new_receiver) as ctor: + client.connect() + + ctor.assert_called_once() + new_receiver.start.assert_called_once() + self.assertIs(client.quote_receiver, new_receiver) + self.assertFalse(client._stop_event.is_set()) + + def test_quote_handler_stops_before_dispatching_rest_of_batch(self): + messages = queue.Queue() + callbacks = [] + fake_client = SimpleNamespace( + logger=Mock(), + quotes=messages, + on_quote=None, + on_trade=None, + ) + handler = EquitiesQuoteHandler(fake_client, bypass_parsing=True) + + def on_trade(item, backlog): + callbacks.append(item) + handler.stop() + + fake_client.on_trade = on_trade + datum = bytearray(49) + datum[0] = 2 + datum[1] = 0 + datum[2] = 24 + datum[25] = 0 + datum[26] = 24 + messages.put(bytes(datum)) + + handler.start() + handler.join(timeout=1) + + self.assertFalse(handler.is_alive()) + self.assertEqual(len(callbacks), 1) + + def test_bypass_parsing_returns_complete_raw_record(self): + for message_type, callback_name in ((0, "on_trade"), (1, "on_quote")): + with self.subTest(message_type=message_type): + callback = Mock() + fake_client = SimpleNamespace( + logger=Mock(), + on_quote=None, + on_trade=None, + ) + setattr(fake_client, callback_name, callback) + handler = EquitiesQuoteHandler(fake_client, bypass_parsing=True) + record = bytearray(range(24)) + record[0] = message_type + record[1] = len(record) + + next_index = handler.parse_message(bytes(record), 0, 7) + + self.assertEqual(next_index, len(record)) + callback.assert_called_once_with(bytes(record), 7) + + def test_quote_handler_ignores_text_and_malformed_binary_messages(self): + client = _make_client() + client.on_trade = Mock() + handler = EquitiesQuoteHandler(client, client.bypass_parsing) + client.quote_handler = handler + handler.start() + client.quotes.put("server text message") + malformed = bytearray(48) + malformed[0] = 2 + malformed[2] = 255 + client.quotes.put(bytes(malformed)) + undersized = bytearray(25) + undersized[0] = 1 + undersized[2] = 1 + client.quotes.put(bytes(undersized)) + + time.sleep(0.05) + + self.assertTrue(handler.is_alive()) + client.on_trade.assert_not_called() + client.disconnect() + + def test_connect_replaces_stopped_quote_handler_and_queue(self): + client = _make_client() + old_handler = client.quote_handler + old_queue = client.quotes + old_handler.stop() + client.refresh_token = Mock() + receiver = Mock() + + with patch.object(EquitiesQuoteHandler, "start") as start, \ + patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", return_value=receiver): + client.connect() + + self.assertIsNot(client.quote_handler, old_handler) + self.assertIsNot(client.quotes, old_queue) + start.assert_called_once() + receiver.start.assert_called_once() + + def test_disconnect_and_reconnect_from_callback_replaces_current_handler(self): + callback_finished = threading.Event() + client_holder = {} + + def on_trade(item, backlog): + client = client_holder["client"] + client.disconnect() + client.connect() + callback_finished.set() + + client = IntrinioRealtimeEquitiesClient( + { + "api_key": "test-key", + "provider": "IEX", + "bypass_parsing": True, + }, + on_trade, + lambda *args: None, + ) + client_holder["client"] = client + original_handler = client.quote_handler + client.refresh_token = Mock() + client.refresh_websocket = Mock() + datum = bytearray(25) + datum[0] = 1 + datum[1] = 0 + datum[2] = 24 + + client.quotes.put(bytes(datum)) + + self.assertTrue(callback_finished.wait(timeout=1)) + original_handler.join(timeout=1) + self.assertFalse(original_handler.is_alive()) + self.assertIsNot(client.quote_handler, original_handler) + client.disconnect() + + def test_refresh_websocket_replaces_disabled_live_receiver(self): + client = _make_client() + old = Mock() + old.is_alive.return_value = True + old.enabled = False + + def join_and_exit(timeout=None): + old.is_alive.return_value = False + + old.join.side_effect = join_and_exit + client.quote_receiver = old + new_receiver = Mock() + + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", return_value=new_receiver): + client.refresh_websocket() + + old.join.assert_called() + new_receiver.start.assert_called_once() + self.assertIs(client.quote_receiver, new_receiver) + + def test_connect_on_healthy_session_does_not_close_socket(self): + client = _make_client() + client.ready = True + client.refresh_token = Mock() + receiver = Mock() + receiver.is_alive.return_value = True + receiver.enabled = True + client.quote_receiver = receiver + ws = Mock() + client.ws = ws + + client.connect() + + ws.close.assert_not_called() + client.refresh_token.assert_not_called() + self.assertIs(client.quote_receiver, receiver) + self.assertTrue(client.ready) + self.assertTrue(receiver.enabled) + + def test_connect_while_receiver_is_starting_does_not_close_socket(self): + client = _make_client() + client.ready = False + client.refresh_token = Mock() + receiver = Mock() + receiver.is_alive.return_value = True + receiver.enabled = True + client.quote_receiver = receiver + ws = Mock() + client.ws = ws + + client.connect() + + ws.close.assert_not_called() + client.refresh_token.assert_not_called() + self.assertIs(client.quote_receiver, receiver) + + def test_concurrent_connects_start_only_one_receiver(self): + client = _make_client() + client.refresh_quote_handler = Mock() + client.refresh_token = Mock() + start_entered = threading.Event() + release_start = threading.Event() + receivers = [] + + class Receiver: + enabled = True + + def __init__(self, owner): + receivers.append(self) + + def start(self): + start_entered.set() + release_start.wait(timeout=1) + + def is_alive(self): + return True + + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", Receiver): + first = threading.Thread(target=client.connect) + second = threading.Thread(target=client.connect) + first.start() + self.assertTrue(start_entered.wait(timeout=1)) + second.start() + time.sleep(0.05) + self.assertEqual(len(receivers), 1) + release_start.set() + first.join(timeout=1) + second.join(timeout=1) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(len(receivers), 1) + + def test_connect_raises_while_disconnect_is_in_progress(self): + client = _make_client() + client.refresh_quote_handler = Mock() + client.refresh_token = Mock() + old_join_entered = threading.Event() + release_disconnect = threading.Event() + + class OldReceiver: + enabled = True + alive = True + join_calls = 0 + + def is_alive(self): + return self.alive + + def join(self, timeout=None): + self.join_calls += 1 + if self.join_calls == 1: + old_join_entered.set() + release_disconnect.wait(timeout=2) + self.alive = False + + old_receiver = OldReceiver() + new_receiver = Mock() + new_receiver.enabled = True + new_receiver.is_alive.return_value = True + client.quote_receiver = old_receiver + client.ws = Mock() + + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", return_value=new_receiver), \ + patch("intriniorealtime.equities_client.close_websocket_app"): + disconnect_thread = threading.Thread(target=client.disconnect) + disconnect_thread.start() + self.assertTrue(old_join_entered.wait(timeout=1)) + + with self.assertRaisesRegex(RuntimeError, "Cannot connect while client is disconnecting"): + client.connect() + release_disconnect.set() + disconnect_thread.join(timeout=1) + + client.connect() + + self.assertFalse(disconnect_thread.is_alive()) + self.assertIs(client.quote_receiver, new_receiver) + + def test_disconnect_clears_websocket_published_late_by_current_receiver(self): + client = _make_client() + late_websocket = Mock() + + class Receiver: + enabled = True + + def join(self, timeout=None): + client.ws = late_websocket + + def is_alive(self): + return False + + receiver = Receiver() + client.quote_receiver = receiver + client.ws = None + + with patch("intriniorealtime.equities_client.close_websocket_app") as abort: + client.disconnect() + + self.assertIsNone(client.quote_receiver) + self.assertIsNone(client.ws) + abort.assert_called_with(late_websocket) + + def test_concurrent_joins_serialize_channel_refreshes(self): + client = _make_client() + client.ready = True + client.ws = Mock() + first_send_entered = threading.Event() + release_first_send = threading.Event() + + def send(_message, _opcode): + if client.ws.send.call_count == 1: + first_send_entered.set() + release_first_send.wait(timeout=1) + + client.ws.send.side_effect = send + first = threading.Thread(target=client.join, args=("AAPL",)) + second = threading.Thread(target=client.join, args=("MSFT",)) + + first.start() + self.assertTrue(first_send_entered.wait(timeout=1)) + second.start() + time.sleep(0.05) + self.assertEqual(client.ws.send.call_count, 1) + release_first_send.set() + first.join(timeout=1) + second.join(timeout=1) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(client.ws.send.call_count, 2) + self.assertEqual(client.joined_channels, {"AAPL", "MSFT"}) + + def test_disconnect_aborts_socket_before_waiting_for_blocked_join(self): + client = _make_client() + client.ready = True + websocket_app = Mock() + client.ws = websocket_app + send_entered = threading.Event() + release_send = threading.Event() + + def send(_message, _opcode): + send_entered.set() + release_send.wait(timeout=1) + + websocket_app.send.side_effect = send + joiner = threading.Thread(target=client.join, args=("AAPL",)) + disconnecter = threading.Thread(target=client.disconnect) + + with patch( + "intriniorealtime.equities_client.close_websocket_app", + side_effect=lambda app: release_send.set(), + ) as abort: + joiner.start() + self.assertTrue(send_entered.wait(timeout=1)) + disconnecter.start() + disconnecter.join(timeout=1) + joiner.join(timeout=1) + + self.assertFalse(disconnecter.is_alive()) + self.assertFalse(joiner.is_alive()) + abort.assert_any_call(websocket_app) + + def test_disconnect_aborts_socket_during_blocked_initial_channel_replay(self): + client = _make_client() + client.channels = {"AAPL"} + client.ws = Mock() + send_entered = threading.Event() + release_send = threading.Event() + + def send(*args, **kwargs): + send_entered.set() + release_send.wait(timeout=1) + + client.ws.send.side_effect = send + opener = threading.Thread(target=client.on_connect) + stopper = threading.Thread(target=client.disconnect) + + with patch( + "intriniorealtime.equities_client.close_websocket_app", + side_effect=lambda app: release_send.set(), + ) as abort: + opener.start() + self.assertTrue(send_entered.wait(timeout=1)) + stopper.start() + stopper.join(timeout=1) + opener.join(timeout=1) + + self.assertFalse(stopper.is_alive()) + self.assertFalse(opener.is_alive()) + abort.assert_called() + self.assertFalse(client.ready) + + def test_partial_channel_refresh_records_successful_sends(self): + client = _make_client() + client.ready = True + client.channels = {"AAPL", "MSFT"} + client.join_binary_message = Mock(side_effect=lambda channel: channel) + client.ws = Mock() + client.ws.send.side_effect = [None, RuntimeError("send failed")] + + with self.assertRaises(RuntimeError): + client.refresh_channels() + + successful_channel = client.ws.send.call_args_list[0].args[0] + failed_channel = client.ws.send.call_args_list[1].args[0] + self.assertEqual(client.joined_channels, {successful_channel}) + + client.ws.send.reset_mock(side_effect=True) + client.refresh_channels() + + client.ws.send.assert_called_once_with( + failed_channel, + websocket.ABNF.OPCODE_BINARY, + ) + self.assertEqual(client.joined_channels, {"AAPL", "MSFT"}) + + client.channels = set() + client.leave_binary_message = Mock(side_effect=lambda channel: channel) + client.ws.send.reset_mock(side_effect=True) + client.ws.send.side_effect = [None, RuntimeError("send failed")] + + with self.assertRaises(RuntimeError): + client.refresh_channels() + + successful_leave = client.ws.send.call_args_list[0].args[0] + failed_leave = client.ws.send.call_args_list[1].args[0] + self.assertNotIn(successful_leave, client.joined_channels) + self.assertEqual(client.joined_channels, {failed_leave}) + + client.ws.send.reset_mock(side_effect=True) + client.refresh_channels() + + client.ws.send.assert_called_once_with( + failed_leave, + websocket.ABNF.OPCODE_BINARY, + ) + self.assertEqual(client.joined_channels, set()) + + def test_connect_queued_before_disconnect_does_not_restart(self): + client = _make_client() + auth_entered = threading.Event() + auth_calls = [] + + def refresh_token(): + auth_calls.append(1) + auth_entered.set() + client._stop_event.wait(timeout=1) + + client.refresh_quote_handler = Mock() + client.refresh_token = Mock(side_effect=refresh_token) + client.refresh_websocket = Mock() + first = threading.Thread(target=client.connect) + second = threading.Thread(target=client.connect) + + first.start() + self.assertTrue(auth_entered.wait(timeout=1)) + second.start() + time.sleep(0.05) + client.disconnect() + first.join(timeout=1) + second.join(timeout=1) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(auth_calls, [1]) + client.refresh_websocket.assert_not_called() + + def test_connect_resets_backoff_on_new_attempt(self): + client = _make_client() + client.last_self_heal_backoff = 4 + client.refresh_token = Mock() + new_receiver = Mock() + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", return_value=new_receiver): + client.connect() + self.assertEqual(client.last_self_heal_backoff, -1) + new_receiver.start.assert_called_once() + + def test_connect_healthy_noop_does_not_reset_backoff(self): + client = _make_client() + client.last_self_heal_backoff = 4 + receiver = Mock() + receiver.is_alive.return_value = True + receiver.enabled = True + client.quote_receiver = receiver + client.connect() + self.assertEqual(client.last_self_heal_backoff, 4) + + def test_receiver_relies_on_websocket_client_fragment_aggregation(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + continuation_callbacks = [] + + def fake_run_forever(app, *args, **kwargs): + continuation_callbacks.append(app.on_cont_message) + receiver.enabled = False + client._stop_event.set() + + with patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.close_websocket_app"): + receiver.run() + + self.assertEqual(continuation_callbacks, [None]) + + def test_refresh_token_passes_timeout(self): + client = _make_client() + with patch("intriniorealtime.equities_client.requests.get") as get: + get.return_value.status_code = 200 + get.return_value.text = "tok" + client.refresh_token() + self.assertEqual(get.call_args.kwargs["timeout"], CONNECT_TIMEOUT_SECONDS) + + def test_connect_does_not_start_receiver_if_stopped_during_auth(self): + client = _make_client() + + def refresh(): + client._stop_event.set() + + client.refresh_token = Mock(side_effect=refresh) + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver") as ctor: + client.connect() + ctor.assert_not_called() + + def test_incomplete_disconnect_blocks_connect_until_receiver_exits(self): + client = _make_client() + release_receiver = threading.Event() + old = threading.Thread(target=release_receiver.wait, daemon=True) + old.enabled = True + old.start() + client.quote_receiver = old + client.ws = Mock() + + ws = client.ws + with patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", -1), \ + patch("intriniorealtime.equities_client.close_websocket_app") as abort: + with self.assertRaisesRegex(RuntimeError, "disconnect incomplete"): + client.disconnect() + + self.assertFalse(old.enabled) + self.assertIs(client.quote_receiver, old) + self.assertIs(client.ws, ws) + self.assertTrue(client._disconnecting) + self.assertGreaterEqual(abort.call_count, 2) + with self.assertRaisesRegex(RuntimeError, "disconnecting"): + client.connect() + + release_receiver.set() + old.join(timeout=1) + deadline = time.time() + 1 + while client._disconnecting and time.time() < deadline: + time.sleep(0.01) + + self.assertFalse(client._disconnecting) + self.assertIsNone(client.quote_receiver) + self.assertIsNone(client.ws) + + def test_deferred_disconnect_waits_for_calling_handler_too(self): + client = _make_client() + release_receiver = threading.Event() + release_handler = threading.Event() + disconnect_returned = threading.Event() + disconnect_errors = [] + + class Receiver(threading.Thread): + def __init__(self): + super().__init__(target=release_receiver.wait, daemon=True) + self.enabled = True + + class Handler(threading.Thread): + def __init__(self): + super().__init__(target=self.run_callback, daemon=True) + self.stopped = False + + def stop(self): + self.stopped = True + + def run_callback(self): + try: + client.disconnect() + except Exception as error: + disconnect_errors.append(error) + disconnect_returned.set() + release_handler.wait() + + receiver = Receiver() + handler = Handler() + receiver.start() + client.quote_receiver = receiver + client.quote_handler = handler + client.ws = Mock() + + with patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", -1): + handler.start() + self.assertTrue(disconnect_returned.wait(timeout=1)) + + self.assertEqual(len(disconnect_errors), 1) + self.assertRegex(str(disconnect_errors[0]), "disconnect incomplete") + self.assertTrue(client._disconnecting) + + release_receiver.set() + receiver.join(timeout=1) + time.sleep(0.05) + self.assertTrue(client._disconnecting) + self.assertTrue(handler.is_alive()) + + release_handler.set() + handler.join(timeout=1) + deadline = time.time() + 1 + while client._disconnecting and time.time() < deadline: + time.sleep(0.01) + + self.assertFalse(client._disconnecting) + self.assertIsNone(client.quote_receiver) + self.assertIsNone(client.quote_handler) + + def test_refresh_websocket_does_not_replace_receiver_still_alive_after_join(self): + client = _make_client() + old = Mock() + old.is_alive.return_value = True + old.enabled = False + old.join = Mock() + client.quote_receiver = old + client.ws = Mock() + + with patch("intriniorealtime.equities_client.close_websocket_app") as abort, \ + patch("intriniorealtime.equities_client.EquitiesQuoteReceiver") as ctor: + with self.assertRaises(RuntimeError): + client.refresh_websocket() + + ctor.assert_not_called() + self.assertIs(client.quote_receiver, old) + self.assertGreaterEqual(abort.call_count, 2) + + def test_open_then_immediate_drop_backs_off(self): + client = _make_client() + client.token = "tok" + client.do_backoff = Mock() + client.refresh_token = Mock() + receiver = EquitiesQuoteReceiver(client) + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] == 1: + app.on_open(app) + return + receiver.enabled = False + client._stop_event.set() + + with patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.equities_client.STABLE_SESSION_SECONDS", 10), \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.close_websocket_app"): + receiver.run() + + self.assertEqual(run_count["n"], 2) + client.do_backoff.assert_called() + + def test_stable_session_drop_skips_first_backoff(self): + client = _make_client() + client.token = "tok" + client.do_backoff = Mock() + client.refresh_token = Mock() + receiver = EquitiesQuoteReceiver(client) + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] == 1: + app.on_open(app) + time.sleep(0.06) + return + receiver.enabled = False + client._stop_event.set() + + with patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.equities_client.STABLE_SESSION_SECONDS", 0.05), \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.close_websocket_app"): + receiver.run() + + self.assertEqual(run_count["n"], 2) + self.assertEqual(client.refresh_token.call_count, 1) + client.do_backoff.assert_not_called() + + def test_completed_handshake_attempt_cancels_its_watchdog(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] >= 2: + receiver.enabled = False + client._stop_event.set() + + def backoff(): + client._stop_event.wait(timeout=0.06) + + client.do_backoff = backoff + client.refresh_token = Mock() + with patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.02), \ + patch("intriniorealtime.equities_client.websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch.object(client.logger, "warning") as warning: + receiver.run() + + timeout_warnings = [ + call for call in warning.call_args_list + if call.args and "connect timed out" in call.args[0] + ] + self.assertEqual(run_count["n"], 2) + self.assertEqual(timeout_warnings, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_options_client.py b/tests/test_options_client.py new file mode 100644 index 0000000..f422ca8 --- /dev/null +++ b/tests/test_options_client.py @@ -0,0 +1,1149 @@ +import logging +import queue +import threading +import time +import unittest +from unittest.mock import Mock, patch + +import websocket + +from intriniorealtime.options_client import ( + Config, + IntrinioRealtimeOptionsClient, + Providers, + _CONNECT_TIMEOUT_SECONDS, + _STABLE_SESSION_SECONDS, + _WebSocket, + _thread_fn, +) + + +def _make_ws(stop_flag, get_token=None, get_url=None, channels=None, data_queue=None, + channels_lock=None, stats=None, stats_lock=None, sent_channels=None, + worker_threads=None, worker_factory=None): + return _WebSocket( + "ws://example/old", + threading.Lock(), + channels_lock or threading.RLock(), + sent_channels if sent_channels is not None else set(), + stats if stats is not None else {"data": 0, "text": 0}, + stats_lock or threading.Lock(), + worker_threads if worker_threads is not None else [], + lambda: channels if channels is not None else set(), + get_token or (lambda: "tok"), + get_url or (lambda token: "ws://example/" + token), + True, + False, + False, + False, + data_queue if data_queue is not None else queue.Queue(), + stop_flag, + worker_factory, + ) + + +def _make_client(): + config = Config( + api_key="test-key", + provider=Providers.MANUAL, + num_threads=1, + manual_ip_address="127.0.0.1", + ) + return IntrinioRealtimeOptionsClient(config, on_trade=lambda trade: None) + + +class OptionsReconnectTests(unittest.TestCase): + def setUp(self): + logging.disable(logging.CRITICAL) + + def tearDown(self): + logging.disable(logging.NOTSET) + + def test_default_stable_session_requires_sustained_health(self): + self.assertGreaterEqual(_STABLE_SESSION_SECONDS, 60) + + def test_config_rejects_non_positive_or_non_integer_worker_counts(self): + for num_threads in (0, -1, True, 1.5, "2"): + with self.subTest(num_threads=num_threads): + with self.assertRaises(ValueError): + Config( + api_key="test-key", + provider=Providers.MANUAL, + num_threads=num_threads, + manual_ip_address="127.0.0.1", + ) + + config = Config( + api_key="test-key", + provider=Providers.MANUAL, + num_threads=1, + manual_ip_address="127.0.0.1", + ) + config.num_threads = 0 + with self.assertRaises(ValueError): + IntrinioRealtimeOptionsClient(config, on_trade=lambda trade: None) + + def test_config_symbol_set_initializes_channels(self): + config = Config( + api_key="test-key", + provider=Providers.MANUAL, + num_threads=1, + manual_ip_address="127.0.0.1", + symbols={"AAPL__220101C00140000", "MSFT"}, + ) + + client = IntrinioRealtimeOptionsClient(config, on_trade=lambda trade: None) + + self.assertEqual( + client._IntrinioRealtimeOptionsClient__channels, + {"AAPL_220101C140.00", "MSFT"}, + ) + + def test_stale_watchdog_does_not_abort_later_attempt(self): + stop = threading.Event() + second_opened = threading.Event() + aborts_after_second_open = [] + run_count = {"n": 0} + + ws = _make_ws(stop) + + def track_abort(app): + if app is ws and second_opened.is_set(): + aborts_after_second_open.append(1) + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] == 1: + return + app.on_open(app) + second_opened.set() + time.sleep(0.2) + stop.set() + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.close_websocket_app", side_effect=track_abort), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertGreaterEqual(run_count["n"], 2) + self.assertEqual(aborts_after_second_open, []) + + def test_reconnect_refreshes_token_without_extra_argument(self): + stop = threading.Event() + + def real_get_token(): + return "tok2" + + get_token = Mock(side_effect=lambda *args, **kwargs: real_get_token(*args, **kwargs)) + get_url = Mock(return_value="ws://example/new") + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] >= 2: + stop.set() + + ws = _make_ws(stop, get_token=get_token, get_url=get_url) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.close_websocket_app"), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + get_token.assert_called_with() + get_url.assert_called_with("tok2") + self.assertEqual(ws.url, "ws://example/new") + + def test_reconnect_does_not_run_forever_with_stale_url_when_token_refresh_fails(self): + stop = threading.Event() + token_calls = {"n": 0} + + def get_token(): + token_calls["n"] += 1 + if token_calls["n"] < 2: + raise RuntimeError("token failed") + return "tok2" + + urls = [] + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + urls.append(app.url) + if run_count["n"] >= 2: + stop.set() + + ws = _make_ws(stop, get_token=get_token, get_url=lambda token: "ws://example/" + token) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.close_websocket_app"), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertEqual(run_count["n"], 2) + self.assertEqual(urls[0], "ws://example/old") + self.assertEqual(urls[1], "ws://example/tok2") + self.assertEqual(token_calls["n"], 2) + + def test_start_after_stop_clears_stop_event_and_connects_again(self): + client = _make_client() + start_calls = [] + + def fake_ws_start(self): + start_calls.append(1) + self.isReady = True + client._stop_event.wait(5) + + with patch("intriniorealtime.options_client.requests.get") as get, \ + patch("intriniorealtime.options_client.time.sleep"), \ + patch.object(_WebSocket, "start", fake_ws_start): + get.return_value.status_code = 200 + get.return_value.text = "tok" + + client.start() + deadline = time.time() + 2 + while not start_calls and time.time() < deadline: + time.sleep(0.01) + self.assertEqual(start_calls, [1]) + + client.stop() + self.assertTrue(client._stop_event.is_set()) + + client.start() + deadline = time.time() + 2 + while len(start_calls) < 2 and time.time() < deadline: + time.sleep(0.01) + self.assertEqual(len(start_calls), 2) + + client.stop() + + def test_websocket_start_is_noop_while_stopped_and_runs_after_clear(self): + stop = threading.Event() + stop.set() + ws = _make_ws(stop) + with patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever") as run_forever, \ + patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("websocket.setdefaulttimeout"): + ws.start() + run_forever.assert_not_called() + + stop.clear() + calls = [] + + def fake_run_forever(app, *args, **kwargs): + calls.append(1) + stop.set() + + with patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.close_websocket_app"): + ws.start() + self.assertEqual(calls, [1]) + + def test_late_open_after_stop_is_aborted_before_becoming_ready(self): + stop = threading.Event() + stop.set() + ws = _make_ws(stop) + + with patch("intriniorealtime.options_client.close_websocket_app") as abort: + ws.on_open(ws) + + self.assertFalse(ws.isReady) + abort.assert_called_once_with(ws) + + def test_late_open_after_timeout_is_aborted_before_becoming_ready(self): + stop = threading.Event() + ws = _make_ws(stop) + + def fake_run_forever(app, *args, **kwargs): + time.sleep(0.08) + app.on_open(app) + stop.set() + + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.03), \ + patch("intriniorealtime.options_client.should_abort_handshake", return_value=False), \ + patch("intriniorealtime.options_client.close_websocket_app") as abort, \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertFalse(ws.isReady) + self.assertTrue(abort.called) + + def test_relies_on_websocket_client_fragment_aggregation(self): + stop = threading.Event() + data_queue = queue.Queue() + ws = _make_ws(stop, data_queue=data_queue) + + self.assertIsNone(ws.on_cont_message) + ws.on_data(ws, b"complete-message", websocket.ABNF.OPCODE_BINARY, True) + self.assertEqual(data_queue.get_nowait(), b"complete-message") + + def test_open_does_not_publish_ready_before_channel_replay_lock(self): + class GateLock: + def __init__(self): + self.entered = threading.Event() + self.release = threading.Event() + + def __enter__(self): + self.entered.set() + self.release.wait(timeout=1) + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + stop = threading.Event() + channels_lock = GateLock() + ws = _make_ws(stop, channels_lock=channels_lock) + opener = threading.Thread(target=ws.on_open, args=(ws,)) + + opener.start() + self.assertTrue(channels_lock.entered.wait(timeout=1)) + self.assertFalse(ws.isReady) + channels_lock.release.set() + opener.join(timeout=1) + + self.assertFalse(opener.is_alive()) + self.assertTrue(ws.isReady) + + def test_open_rechecks_stop_before_publishing_ready(self): + class GateLock: + def __init__(self): + self.entered = threading.Event() + self.release = threading.Event() + + def __enter__(self): + self.entered.set() + self.release.wait(timeout=1) + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + stop = threading.Event() + channels_lock = GateLock() + ws = _make_ws(stop, channels_lock=channels_lock) + opener = threading.Thread(target=ws.on_open, args=(ws,)) + + opener.start() + self.assertTrue(channels_lock.entered.wait(timeout=1)) + stop.set() + channels_lock.release.set() + opener.join(timeout=1) + + self.assertFalse(opener.is_alive()) + self.assertFalse(ws.isReady) + + def test_open_replaces_a_terminated_worker_before_publishing_ready(self): + stop = threading.Event() + terminated = threading.Thread(target=lambda: None) + terminated.start() + terminated.join(timeout=1) + replacements = [] + + def make_worker(_index): + worker = threading.Thread(target=stop.wait) + replacements.append(worker) + return worker + + workers = [terminated] + ws = _make_ws( + stop, + worker_threads=workers, + worker_factory=make_worker, + ) + + ws.on_open(ws) + + self.assertTrue(ws.isReady) + self.assertIs(workers[0], replacements[0]) + self.assertTrue(replacements[0].is_alive()) + stop.set() + replacements[0].join(timeout=1) + + def test_open_rolls_back_ready_when_worker_start_fails(self): + stop = threading.Event() + terminated = threading.Thread(target=lambda: None) + terminated.start() + terminated.join(timeout=1) + ws = _make_ws(stop, worker_threads=[terminated]) + + with patch("intriniorealtime.options_client.close_websocket_app") as close: + with self.assertRaisesRegex(RuntimeError, "cannot be restarted"): + ws.on_open(ws) + + self.assertFalse(ws.isReady) + close.assert_called_once_with(ws) + + def test_open_aborts_connection_when_channel_replay_fails(self): + stop = threading.Event() + sent_channels = set() + ws = _make_ws( + stop, + channels={"AAPL"}, + sent_channels=sent_channels, + ) + ws.send_binary = Mock(side_effect=RuntimeError("send failed")) + + with patch("intriniorealtime.options_client.close_websocket_app") as abort: + with self.assertRaises(RuntimeError): + ws.on_open(ws) + + self.assertFalse(ws.isReady) + self.assertEqual(sent_channels, set()) + abort.assert_called_once_with(ws) + + def test_get_channels_returns_snapshot(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} + + snapshot = client._IntrinioRealtimeOptionsClient__get_channels() + snapshot.add("MSFT") + + self.assertEqual(client._IntrinioRealtimeOptionsClient__channels, {"AAPL"}) + + def test_stats_are_isolated_between_clients(self): + first = _make_client() + second = _make_client() + first_data = first._IntrinioRealtimeOptionsClient__data + first_stats = first._IntrinioRealtimeOptionsClient__stats + first_stats_lock = first._IntrinioRealtimeOptionsClient__stats_lock + ws = _make_ws( + threading.Event(), + data_queue=first_data, + stats=first_stats, + stats_lock=first_stats_lock, + ) + + ws.on_data(ws, b"message", websocket.ABNF.OPCODE_BINARY, True) + + self.assertEqual(first.get_stats(), (1, 0, 1)) + self.assertEqual(second.get_stats(), (0, 0, 0)) + + def test_worker_stops_before_dispatching_rest_of_batch(self): + data = queue.Queue() + stop = threading.Event() + callbacks = [] + + def on_quote(quote): + callbacks.append(quote) + stop.set() + + datum = bytearray(105) + datum[0] = 2 + datum[1 + 22] = 1 + datum[1 + 52 + 22] = 1 + data.put(bytes(datum)) + + worker = threading.Thread( + target=_thread_fn, + args=(0, data, None, on_quote, None, None, stop), + ) + worker.start() + worker.join(timeout=1) + + self.assertFalse(worker.is_alive()) + self.assertEqual(len(callbacks), 1) + + def test_truncated_message_does_not_terminate_worker(self): + data = queue.Queue() + stop = threading.Event() + worker = threading.Thread( + target=_thread_fn, + args=(0, data, None, None, None, None, stop), + ) + worker.start() + + data.put(b"\x01") + time.sleep(0.05) + + self.assertTrue(worker.is_alive()) + stop.set() + worker.join(timeout=2) + self.assertFalse(worker.is_alive()) + + def test_stop_preserves_channels_and_clears_started(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} + client._IntrinioRealtimeOptionsClient__sent_channels = {"AAPL"} + client._IntrinioRealtimeOptionsClient__is_started = True + ws = Mock() + ws.isReady = True + client._IntrinioRealtimeOptionsClient__webSocket = ws + client._IntrinioRealtimeOptionsClient__socket_thread = None + + client.stop() + + self.assertEqual(client._IntrinioRealtimeOptionsClient__channels, {"AAPL"}) + self.assertEqual(client._IntrinioRealtimeOptionsClient__sent_channels, set()) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + ws.send_binary.assert_not_called() + ws.stop.assert_called() + + def test_join_after_stop_does_not_wait_for_ready(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__is_started = True + client._IntrinioRealtimeOptionsClient__webSocket = None + with patch("intriniorealtime.options_client.time.sleep"): + client.stop() + + with patch.object(client, "_IntrinioRealtimeOptionsClient__all_ready") as ready: + client.join("MSFT") + ready.assert_not_called() + self.assertIn("MSFT", client._IntrinioRealtimeOptionsClient__channels) + + def test_join_does_not_wait_while_connection_is_unavailable(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__is_started = True + client._IntrinioRealtimeOptionsClient__webSocket = None + + with patch.object(client, "_IntrinioRealtimeOptionsClient__all_ready") as ready: + client.join("MSFT") + client.join_firehose() + + ready.assert_not_called() + self.assertIn("MSFT", client._IntrinioRealtimeOptionsClient__channels) + self.assertIn("$FIREHOSE", client._IntrinioRealtimeOptionsClient__channels) + + def test_reconnect_token_refresh_invalidates_cached_token(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__token = ("cached", time.time()) + + with patch("intriniorealtime.options_client.requests.get") as get: + get.return_value.status_code = 200 + get.return_value.text = "fresh" + token = client._IntrinioRealtimeOptionsClient__refresh_token() + + self.assertEqual(token, "fresh") + get.assert_called_once() + + def test_leave_drops_channel_if_send_raises(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} + ws = Mock() + ws.isReady = True + ws.send_binary.side_effect = RuntimeError("connection closed") + client._IntrinioRealtimeOptionsClient__webSocket = ws + + with self.assertRaises(RuntimeError): + client.leave("AAPL") + + self.assertNotIn("AAPL", client._IntrinioRealtimeOptionsClient__channels) + + def test_join_retries_send_after_previous_send_raises(self): + client = _make_client() + ws = Mock() + ws.isReady = True + ws.send_binary.side_effect = RuntimeError("connection closed") + client._IntrinioRealtimeOptionsClient__webSocket = ws + + with self.assertRaises(RuntimeError): + client.join("AAPL") + + ws.send_binary.side_effect = None + client.join("AAPL") + + self.assertIn("AAPL", client._IntrinioRealtimeOptionsClient__channels) + self.assertEqual(ws.send_binary.call_count, 2) + + def test_stop_joins_socket_thread_with_timeout(self): + client = _make_client() + thread = Mock() + thread.ident = 1 + thread.is_alive.return_value = False + client._IntrinioRealtimeOptionsClient__socket_thread = thread + client._IntrinioRealtimeOptionsClient__webSocket = None + with patch("intriniorealtime.options_client.time.sleep"): + client.stop() + timeout = thread.join.call_args.kwargs["timeout"] + self.assertGreater(timeout, 0) + self.assertLessEqual(timeout, _CONNECT_TIMEOUT_SECONDS + 1) + + def test_stop_from_worker_callback_finishes_after_worker_exits(self): + client = _make_client() + callback_entered = threading.Event() + stop_returned = threading.Event() + release_callback = threading.Event() + + def callback_worker(): + callback_entered.set() + client.stop() + stop_returned.set() + release_callback.wait(timeout=1) + + worker = threading.Thread(target=callback_worker) + client._IntrinioRealtimeOptionsClient__worker_threads = [worker] + client._IntrinioRealtimeOptionsClient__socket_thread = None + client._IntrinioRealtimeOptionsClient__webSocket = None + client._IntrinioRealtimeOptionsClient__is_started = True + worker.start() + + self.assertTrue(callback_entered.wait(timeout=1)) + self.assertTrue(stop_returned.wait(timeout=1)) + self.assertTrue(client._IntrinioRealtimeOptionsClient__is_stopping) + + waiting_stop = threading.Thread(target=client.stop) + waiting_stop.start() + waiting_stop.join(timeout=0.05) + self.assertTrue(waiting_stop.is_alive()) + + with self.assertRaisesRegex(RuntimeError, "client is stopping"): + client.start() + + release_callback.set() + worker.join(timeout=1) + waiting_stop.join(timeout=1) + + self.assertTrue(client._stop_event.is_set()) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_stopping) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + self.assertFalse(waiting_stop.is_alive()) + + def test_concurrent_worker_stops_do_not_wait_on_each_other(self): + client = _make_client() + barrier = threading.Barrier(3) + results = [] + + def callback_worker(): + barrier.wait() + try: + client.stop() + results.append(None) + except Exception as error: + results.append(error) + + workers = [threading.Thread(target=callback_worker) for _ in range(2)] + client._IntrinioRealtimeOptionsClient__worker_threads = workers + client._IntrinioRealtimeOptionsClient__socket_thread = None + client._IntrinioRealtimeOptionsClient__is_started = True + for worker in workers: + worker.start() + + barrier.wait() + for worker in workers: + worker.join(timeout=2) + + deadline = time.time() + 1 + while client._IntrinioRealtimeOptionsClient__is_stopping and time.time() < deadline: + time.sleep(0.01) + + self.assertEqual(results, [None, None]) + self.assertFalse(any(worker.is_alive() for worker in workers)) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_stopping) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + + def test_open_then_drop_twice_backs_off_on_second_reconnect(self): + stop = threading.Event() + waits = [] + original_wait = threading.Event.wait + + def tracking_wait(self, timeout=None): + waits.append(timeout) + return original_wait(self, timeout) + + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + app.on_open(app) + if run_count["n"] >= 3: + stop.set() + + ws = _make_ws(stop) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._STABLE_SESSION_SECONDS", 10), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0.01]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.close_websocket_app"), \ + patch("threading.Event.wait", tracking_wait), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertEqual(run_count["n"], 3) + self.assertIn(0.01, waits) + + def test_stable_session_drop_skips_first_backoff(self): + stop = threading.Event() + waits = [] + original_wait = threading.Event.wait + + def tracking_wait(self, timeout=None): + waits.append(timeout) + return original_wait(self, timeout) + + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] == 1: + app.on_open(app) + time.sleep(0.06) + return + stop.set() + + ws = _make_ws(stop) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._STABLE_SESSION_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [10]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.close_websocket_app"), \ + patch("threading.Event.wait", tracking_wait), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertEqual(run_count["n"], 2) + self.assertNotIn(10, waits) + + def test_failed_handshake_still_backs_off_before_reconnect(self): + stop = threading.Event() + waits = [] + original_wait = threading.Event.wait + + def tracking_wait(self, timeout=None): + waits.append(timeout) + return original_wait(self, timeout) + + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] == 1: + return + stop.set() + + ws = _make_ws(stop) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0.01]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.close_websocket_app"), \ + patch("threading.Event.wait", tracking_wait), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertEqual(run_count["n"], 2) + self.assertIn(0.01, waits) + + def test_completed_handshake_attempt_cancels_its_watchdog(self): + stop = threading.Event() + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] >= 2: + stop.set() + + ws = _make_ws(stop) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.02), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0.06]), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.options_client._log.warning") as warning: + ws.start() + + timeout_warnings = [ + call for call in warning.call_args_list + if call.args and "Connect timed out" in call.args[0] + ] + self.assertEqual(run_count["n"], 2) + self.assertEqual(timeout_warnings, []) + + def test_stop_sets_flag_before_socket_abort(self): + client = _make_client() + saw_stop = {} + ws = Mock() + ws.isReady = True + client._IntrinioRealtimeOptionsClient__webSocket = ws + + def stop(): + saw_stop["set"] = client._stop_event.is_set() + + ws.stop.side_effect = stop + client.stop() + self.assertTrue(saw_stop["set"]) + + def test_stop_aborts_socket_before_waiting_for_blocked_join(self): + client = _make_client() + ws = Mock() + ws.isReady = True + send_entered = threading.Event() + release_send = threading.Event() + + def send_binary(_message): + send_entered.set() + release_send.wait(timeout=1) + + ws.send_binary.side_effect = send_binary + ws.stop.side_effect = release_send.set + client._IntrinioRealtimeOptionsClient__webSocket = ws + joiner = threading.Thread(target=client.join, args=("AAPL",)) + stopper = threading.Thread(target=client.stop) + + joiner.start() + self.assertTrue(send_entered.wait(timeout=1)) + stopper.start() + stopper.join(timeout=1) + joiner.join(timeout=1) + + self.assertFalse(stopper.is_alive()) + self.assertFalse(joiner.is_alive()) + ws.stop.assert_called() + + def test_start_raises_if_stopped_during_auth(self): + client = _make_client() + + def get_token(): + client._stop_event.set() + return None + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", side_effect=get_token), \ + patch.object(threading.Thread, "start") as thread_start: + with self.assertRaisesRegex(RuntimeError, "Cannot start while client is stopping"): + client.start() + + thread_start.assert_not_called() + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + + def test_start_retries_until_token_then_connects(self): + client = _make_client() + tokens = {"n": 0} + start_calls = [] + + def get_token(): + tokens["n"] += 1 + if tokens["n"] < 3: + return None + return "tok" + + def fake_ws_start(self): + start_calls.append(1) + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", side_effect=get_token), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ + patch.object(_WebSocket, "start", fake_ws_start): + client.start() + + self.assertEqual(tokens["n"], 3) + self.assertEqual(start_calls, [1]) + self.assertTrue(client._IntrinioRealtimeOptionsClient__is_started) + + def test_stop_during_socket_thread_start_raises_lifecycle_error(self): + client = _make_client() + + def stop_before_thread_runs(): + client.stop() + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", return_value="tok"), \ + patch.object(threading.Thread, "start", side_effect=stop_before_thread_runs), \ + patch("intriniorealtime.options_client.time.sleep"): + with self.assertRaisesRegex(RuntimeError, "Cannot start while client is stopping"): + client.start() + + self.assertTrue(client._stop_event.is_set()) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + + def test_start_replaces_data_queue_from_previous_session(self): + client = _make_client() + old_queue = client._IntrinioRealtimeOptionsClient__data + old_queue.put(b"stale") + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", return_value="tok"), \ + patch.object(_WebSocket, "start"): + client.start() + + new_queue = client._IntrinioRealtimeOptionsClient__data + self.assertIsNot(new_queue, old_queue) + self.assertTrue(new_queue.empty()) + self.assertEqual(old_queue.get_nowait(), b"stale") + + def test_start_on_healthy_session_is_noop(self): + client = _make_client() + thread = Mock() + thread.is_alive.return_value = True + client._IntrinioRealtimeOptionsClient__socket_thread = thread + client._IntrinioRealtimeOptionsClient__is_started = True + client._stop_event.clear() + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token") as get_token: + client.start() + + get_token.assert_not_called() + + def test_start_waits_for_previous_socket_thread_then_starts(self): + client = _make_client() + leftover = Mock() + leftover.ident = 1 + leftover.is_alive.side_effect = [True, False] + leftover.join = Mock() + client._IntrinioRealtimeOptionsClient__socket_thread = leftover + client._IntrinioRealtimeOptionsClient__is_started = False + client._stop_event.set() + start_calls = [] + + def fake_ws_start(self): + start_calls.append(1) + + with patch("intriniorealtime.options_client.requests.get") as get, \ + patch.object(_WebSocket, "start", fake_ws_start): + get.return_value.status_code = 200 + get.return_value.text = "tok" + client.start() + + leftover.join.assert_called() + self.assertEqual(start_calls, [1]) + + def test_start_fails_when_previous_socket_thread_does_not_exit(self): + client = _make_client() + leftover = Mock() + leftover.ident = 1 + leftover.is_alive.return_value = True + client._IntrinioRealtimeOptionsClient__socket_thread = leftover + client._IntrinioRealtimeOptionsClient__is_started = False + client._stop_event.set() + + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0), \ + patch("intriniorealtime.options_client.close_websocket_app"): + with self.assertRaisesRegex(RuntimeError, "Previous socket thread"): + client.start() + + leftover.join.assert_called_once_with(timeout=1) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_starting) + + def test_start_raises_clear_error_while_stop_is_in_progress(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__is_stopping = True + + with self.assertRaisesRegex(RuntimeError, "Cannot start while client is stopping"): + client.start() + + def test_concurrent_start_waits_and_reports_active_start_failure(self): + client = _make_client() + token_entered = threading.Event() + release_token = threading.Event() + errors = [] + + def fail_token(): + token_entered.set() + release_token.wait(timeout=1) + raise RuntimeError("auth failed") + + def start_client(): + try: + client.start() + except Exception as error: + errors.append(str(error)) + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", side_effect=fail_token): + first = threading.Thread(target=start_client) + second = threading.Thread(target=start_client) + first.start() + self.assertTrue(token_entered.wait(timeout=1)) + second.start() + second.join(timeout=0.05) + self.assertTrue(second.is_alive()) + + release_token.set() + first.join(timeout=1) + second.join(timeout=1) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(len(errors), 2) + self.assertTrue(any(error == "auth failed" for error in errors)) + self.assertTrue(any(error.startswith("Concurrent start failed: auth failed") for error in errors)) + + def test_stop_completes_teardown_with_active_socket(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} + client._IntrinioRealtimeOptionsClient__is_started = True + ws = Mock() + ws.isReady = True + client._IntrinioRealtimeOptionsClient__webSocket = ws + thread = Mock() + thread.ident = 1 + thread.is_alive.return_value = False + client._IntrinioRealtimeOptionsClient__socket_thread = thread + + client.stop() + + ws.stop.assert_called() + thread.join.assert_called() + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + + def test_start_uses_new_stop_event_so_old_workers_stay_stopped(self): + client = _make_client() + old_event = client._stop_event + old_event.set() + old_worker = Mock() + old_worker.ident = 1 + old_worker.is_alive.return_value = False + client._IntrinioRealtimeOptionsClient__worker_threads = [old_worker] + client._IntrinioRealtimeOptionsClient__is_started = False + start_calls = [] + + def fake_ws_start(self): + start_calls.append(1) + + with patch("intriniorealtime.options_client.requests.get") as get, \ + patch.object(_WebSocket, "start", fake_ws_start): + get.return_value.status_code = 200 + get.return_value.text = "tok" + client.start() + + self.assertIsNot(client._stop_event, old_event) + self.assertTrue(old_event.is_set()) + self.assertFalse(client._stop_event.is_set()) + self.assertNotIn(old_worker, client._IntrinioRealtimeOptionsClient__worker_threads) + self.assertEqual(start_calls, [1]) + + def test_stop_reports_workers_that_do_not_exit(self): + client = _make_client() + worker = Mock() + worker.ident = 1 + worker.is_alive.return_value = True + client._IntrinioRealtimeOptionsClient__worker_threads = [worker] + client._IntrinioRealtimeOptionsClient__socket_thread = None + + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0): + with self.assertRaisesRegex(RuntimeError, "stop incomplete"): + client.stop() + + self.assertTrue(client._IntrinioRealtimeOptionsClient__is_started) + with self.assertRaisesRegex(RuntimeError, "Previous worker threads"): + client.start() + + def test_stop_rechecks_survivors_after_all_joins(self): + client = _make_client() + worker = Mock() + worker.ident = 1 + worker.is_alive.side_effect = [True, False] + client._IntrinioRealtimeOptionsClient__worker_threads = [worker] + client._IntrinioRealtimeOptionsClient__socket_thread = None + + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0): + client.stop() + + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + + def test_concurrent_stop_waits_for_active_stop(self): + client = _make_client() + join_entered = threading.Event() + release_join = threading.Event() + worker_alive = {"value": True} + worker = Mock() + worker.ident = 1 + worker.is_alive.side_effect = lambda: worker_alive["value"] + + def join(timeout=None): + join_entered.set() + release_join.wait(timeout=1) + worker_alive["value"] = False + + worker.join.side_effect = join + client._IntrinioRealtimeOptionsClient__worker_threads = [worker] + client._IntrinioRealtimeOptionsClient__socket_thread = None + failures = [] + + def stop_client(): + try: + client.stop() + except Exception as error: + failures.append(error) + + first = threading.Thread(target=stop_client) + second = threading.Thread(target=stop_client) + first.start() + self.assertTrue(join_entered.wait(timeout=1)) + second.start() + second.join(timeout=0.05) + self.assertTrue(second.is_alive()) + + release_join.set() + first.join(timeout=1) + second.join(timeout=1) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(failures, []) + + def test_stop_uses_one_deadline_for_all_thread_joins(self): + client = _make_client() + clock = {"now": 100.0} + workers = [] + + for _ in range(2): + worker = Mock() + worker.ident = 1 + worker.is_alive.return_value = False + + def join(timeout=None, clock=clock): + clock["now"] += 0.6 + + worker.join.side_effect = join + workers.append(worker) + + client._IntrinioRealtimeOptionsClient__worker_threads = workers + client._IntrinioRealtimeOptionsClient__socket_thread = None + + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0), \ + patch("intriniorealtime.options_client.time.monotonic", side_effect=lambda: clock["now"]): + client.stop() + + first_timeout = workers[0].join.call_args.kwargs["timeout"] + second_timeout = workers[1].join.call_args.kwargs["timeout"] + self.assertEqual(first_timeout, 1.0) + self.assertAlmostEqual(second_timeout, 0.4) + + def test_start_aborts_leftover_websocket_while_joining(self): + client = _make_client() + leftover = Mock() + leftover.ident = 1 + leftover.is_alive.side_effect = [True, False] + leftover.join = Mock() + ws = Mock() + client._IntrinioRealtimeOptionsClient__socket_thread = leftover + client._IntrinioRealtimeOptionsClient__webSocket = ws + client._IntrinioRealtimeOptionsClient__is_started = False + client._stop_event.set() + start_calls = [] + + def fake_ws_start(self): + start_calls.append(1) + + with patch("intriniorealtime.options_client.close_websocket_app") as abort, \ + patch("intriniorealtime.options_client.requests.get") as get, \ + patch.object(_WebSocket, "start", fake_ws_start): + get.return_value.status_code = 200 + get.return_value.text = "tok" + client.start() + + abort.assert_called() + leftover.join.assert_called() + self.assertEqual(start_calls, [1]) + + def test_start_raises_if_stop_occurs_during_leftover_join(self): + client = _make_client() + leftover = Mock() + leftover.ident = 1 + leftover.is_alive.return_value = True + + def join(timeout=None): + client._stop_generation += 1 + client._stop_event.set() + + leftover.join.side_effect = join + client._IntrinioRealtimeOptionsClient__socket_thread = leftover + client._IntrinioRealtimeOptionsClient__is_started = False + client._stop_event.set() + + with patch("intriniorealtime.options_client.close_websocket_app"), \ + patch.object(client, "_IntrinioRealtimeOptionsClient__get_token") as get_token, \ + patch.object(threading.Thread, "start") as thread_start: + with self.assertRaisesRegex(RuntimeError, "Cannot start while client is stopping"): + client.start() + + get_token.assert_not_called() + thread_start.assert_not_called() + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) diff --git a/tests/test_websocket_abort.py b/tests/test_websocket_abort.py new file mode 100644 index 0000000..be232d5 --- /dev/null +++ b/tests/test_websocket_abort.py @@ -0,0 +1,75 @@ +import threading +import time +import unittest +from unittest.mock import Mock + +from intriniorealtime._websocket import close_websocket_app, should_abort_handshake + + +class CloseWebsocketAppTests(unittest.TestCase): + def test_uses_only_public_close_method(self): + app = Mock(spec=["close"]) + + close_websocket_app(app) + + app.close.assert_called_once_with() + + def test_noop_on_none(self): + close_websocket_app(None) + + def test_swallows_close_error(self): + app = Mock(spec=["close"]) + app.close.side_effect = RuntimeError("close failed") + close_websocket_app(app) + + +class ShouldAbortHandshakeTests(unittest.TestCase): + def test_true_on_timeout_for_current_generation(self): + handshake = threading.Event() + stop = threading.Event() + self.assertTrue(should_abort_handshake( + handshake, stop, 1, 1, time.monotonic() + 0.01 + )) + + def test_false_when_handshake_completes(self): + handshake = threading.Event() + handshake.set() + stop = threading.Event() + self.assertFalse(should_abort_handshake( + handshake, stop, 1, 1, time.monotonic() + 0.01 + )) + + def test_false_for_stale_generation(self): + handshake = threading.Event() + stop = threading.Event() + self.assertFalse(should_abort_handshake( + handshake, stop, 1, 2, time.monotonic() + 0.01 + )) + + def test_false_immediately_when_stopped(self): + handshake = threading.Event() + stop = threading.Event() + stop.set() + start = time.monotonic() + self.assertFalse(should_abort_handshake( + handshake, stop, 1, 1, time.monotonic() + 1.0 + )) + self.assertLess(time.monotonic() - start, 0.1) + + def test_callable_current_generation_is_read_after_wait(self): + handshake = threading.Event() + stop = threading.Event() + state = {"g": 1} + + def bump(): + time.sleep(0.02) + state["g"] = 2 + + threading.Thread(target=bump, daemon=True).start() + self.assertFalse(should_abort_handshake( + handshake, stop, 1, lambda: state["g"], time.monotonic() + 0.05 + )) + + +if __name__ == "__main__": + unittest.main()