diff --git a/bot/signal.py b/bot/signal.py index 6881123..dd2265a 100644 --- a/bot/signal.py +++ b/bot/signal.py @@ -1,5 +1,5 @@ -from typing import Optional -import datetime, websocket, json, copy, queue, threading +from typing import Optional, Callable, Union +import datetime, websocket, json, copy, queue, threading, time from . import exceptions class Signal: @@ -9,13 +9,6 @@ class Signal: only with `/signals` """ - # As of now signals should be kept internal until initial Python SDK - # scope is defined. Then the SignalType from the tests can be reused. - available_signals = [ - 'messages.new', 'message.delivered', 'node.ready', - 'node.started', 'node.login', 'node.stopped' - ] - def __init__(self, url: str): self.__url = url self.__data = {} @@ -24,6 +17,16 @@ class Signal: # For real time data extraction self.__queue = queue.Queue() self.__thread = None + # State for the long-lived `connect()` + `expect()` flow. These + # attributes are intentionally single-underscore so `SignalExpectation` + # can read them without name-mangling gymnastics. + self._cond = threading.Condition() + self._received_by_type: dict[str, list[dict]] = {} + self._received_all: list[tuple[int, str, dict]] = [] + self._seq = 0 + self.__connect_thread: Optional[threading.Thread] = None + self.__connect_ws: Optional[websocket.WebSocketApp] = None + self.__should_stop = False def __on_open(self, ws: websocket.WebSocketApp): """ @@ -148,3 +151,217 @@ class Signal: yield data + def connect(self): + """ + Open a persistent websocket connection that buffers every incoming signal + for inspection via `expect()`. Must be called before `expect()`. Idempotent + - calling twice while a connection is alive is a no-op. + + The connection auto-reconnects on disconnect with exponential backoff, + so `expect()` keeps working across transient backend hiccups. + """ + if self.__connect_thread and self.__connect_thread.is_alive(): + return + self.__should_stop = False + self.__connect_thread = threading.Thread(target=self.__connect_loop, daemon=True) + self.__connect_thread.start() + + def disconnect(self): + """ + Tear down the persistent connection opened by `connect()`. Idempotent. + Safe to call even if `connect()` was never invoked. + """ + self.__should_stop = True + if self.__connect_ws is not None: + self.__connect_ws.close() + if self.__connect_thread and self.__connect_thread.is_alive(): + self.__connect_thread.join(timeout=1) + + def __connect_loop(self): + """ + Background loop that maintains the persistent websocket and reconnects + with exponential backoff on disconnect, until `disconnect()` is called. + """ + retry_delay = 0.5 + max_delay = 5.0 + while not self.__should_stop: + self.__connect_ws = websocket.WebSocketApp( + self.__url, + on_message=self.__buffer_message, + on_error=self.__on_error, + ) + self.__connect_ws.run_forever() + if self.__should_stop: + break + time.sleep(retry_delay) + retry_delay = min(retry_delay * 2, max_delay) + + def __buffer_message(self, ws: websocket.WebSocketApp, raw: str): + """ + Persistent-connection on_message handler. Parses the raw signal, normalizes + the payload, and appends it to both the per-type buffer and the global + ordered stream. Notifies any threads waiting in `SignalExpectation.__exit__`. + """ + try: + signal: dict = json.loads(raw) + except json.JSONDecodeError: + return + + signal_type = signal.get("type") + if not signal_type: + return + + event: dict = signal.get("event") or {} + data = { + "type": signal_type, + "timestamp": datetime.datetime.fromtimestamp(signal["timestamp"]), + "is_error": event.get("error") is not None, + "error_message": event.get("error"), + "event": event, + } + with self._cond: + self._seq += 1 + self._received_by_type.setdefault(signal_type, []).append(data) + self._received_all.append((self._seq, signal_type, data)) + self._cond.notify_all() + + def expect( + self, + signal_type: str, + *, + count: int = 1, + accept_fn: Optional[Callable[[dict], bool]] = None, + pattern: Optional[str] = None, + predicate: Optional[Callable[[dict], bool]] = None, + timeout: float = 20.0, + start: Union[str, int] = "now", + ) -> "SignalExpectation": + """ + Return a context manager that waits for `count` matching signals of `signal_type` + to arrive after entering the `with` block. Perform the triggering action inside the + block. After the block exits the matched signals are exposed as `exp.result` (or + `exp.results` when `count > 1`). + + Requires `connect()` to be running. Raises `TimeoutError` if the expected signals + do not arrive within `timeout` seconds. + + Parameters: + - `signal_type` - the "type" as it appears in Status Backend + - `count` - number of matching signals to wait for (default 1) + - `accept_fn` - optional filter; called with the buffered signal dict + - `pattern` - optional substring to match in the JSON-serialized signal + - `predicate` - alias for `accept_fn`. Only one of these three may be set. + - `timeout` - max seconds to wait inside `__exit__` + - `start` - `"now"` (default, only future signals), `"beginning"` (search + from the start of the buffer) or an explicit int index + """ + if not self.__connect_thread or not self.__connect_thread.is_alive(): + raise exceptions.SignalError("connect() must be called before expect()") + + return SignalExpectation( + self, + signal_type, + count=count, + accept_fn=accept_fn, + pattern=pattern, + predicate=predicate, + timeout=timeout, + start=start, + ) + + +class SignalExpectation: + """ + Context manager that asserts an action triggers one or more matching signals. + + On `__enter__` it snapshots the current size of the per-type signal buffer so + any signals received before the action are ignored (race-safe). On `__exit__` + it blocks until `count` matching signals appear past the snapshot or until + `timeout` expires, in which case `TimeoutError` is raised. + + Typical usage: + + signal.connect() + with signal.expect("messages.new") as exp: + account.send_message(chat_id, "hello") + print(exp.result) + """ + + def __init__( + self, + signal: Signal, + signal_type: str, + *, + count: int = 1, + accept_fn: Optional[Callable[[dict], bool]] = None, + pattern: Optional[str] = None, + predicate: Optional[Callable[[dict], bool]] = None, + timeout: float = 20.0, + start: Union[str, int] = "now", + ): + if count < 1: + raise ValueError("count must be >= 1") + + filters_set = sum(1 for v in (accept_fn, pattern, predicate) if v is not None) + if filters_set > 1: + raise ValueError("Only one of accept_fn, pattern, predicate can be specified") + + self._signal = signal + self._signal_type = signal_type + self._count = count + self._timeout = float(timeout) + self._start = start + self._start_index = 0 + + if pattern is not None: + self._accept_fn: Optional[Callable[[dict], bool]] = lambda s: pattern in json.dumps(s, default=str) + elif predicate is not None: + self._accept_fn = predicate + else: + self._accept_fn = accept_fn + + self.result: Optional[Union[dict, list[dict]]] = None + self.results: Optional[list[dict]] = None + + def __enter__(self) -> "SignalExpectation": + with self._signal._cond: + buffer = self._signal._received_by_type.setdefault(self._signal_type, []) + if self._start == "now": + self._start_index = len(buffer) + elif self._start == "beginning": + self._start_index = 0 + elif isinstance(self._start, int): + if self._start < 0: + raise ValueError("start index must be >= 0") + self._start_index = self._start + else: + raise ValueError(f"Unsupported start mode: {self._start!r}") + return self + + def __exit__(self, exc_type, exc, tb): + # Don't swallow exceptions raised inside the `with` body. + if exc_type is not None: + return False + + deadline = time.time() + self._timeout + with self._signal._cond: + while True: + buffer = self._signal._received_by_type.get(self._signal_type, []) + candidates = buffer[self._start_index:] + if self._accept_fn is not None: + candidates = [s for s in candidates if self._accept_fn(s)] + + if len(candidates) >= self._count: + self.results = candidates[: self._count] + self.result = self.results[0] if self._count == 1 else self.results + return False + + remaining = deadline - time.time() + if remaining <= 0: + raise TimeoutError( + f"Expected {self._count} signal(s) of type {self._signal_type!r}, " + f"got {len(candidates)} within {self._timeout}s" + ) + + self._signal._cond.wait(timeout=remaining) + diff --git a/docs/account.md b/docs/account.md index 0abecf7..6d41abe 100644 --- a/docs/account.md +++ b/docs/account.md @@ -1026,19 +1026,38 @@ When a new profile picture is set, any previous image in the **assets** folder i ### `signal` -The property exists in `Account` because signals require an **active logged‑in session**. Attempting to use signals before calling `login()` will raise an exception. Signals are low‑level events emitted by the Status Backend. Examples include: +The property exists in `Account` because signals require an **active logged‑in session**. Attempting to use signals before calling `login()` will raise an exception. Signals are low‑level events emitted by the Status Backend. -- `messages.new` -- `message.delivered` -- `node.ready` -- `node.started` -- `node.login` -- `node.stopped` - -The property exposes two primary methods: +The property exposes the following methods: - `signal.get()` - fetch a single event. If the event is not found, you may end up in an infinite loop. - `signal.listen()` - stream events continuously. Example usage of this is found in [`listen_messages()`](./account.md#listen_messages) +- `signal.connect()` - open a persistent websocket connection that buffers **every** incoming signal in the background. Must be called before `expect()`. +- `signal.disconnect()` - tear down the persistent connection opened by `connect()`. +- `signal.expect()` - return a context manager that waits for one or more matching signals to arrive **after** you perform an action. This is the recommended way to make **async message calls**, since it removes the race conditions and infinite-loop risk of `get()`. + +```python +from bot import Account + +account = Account() +params = { + "name": "status-app-bot", + "password": "SNTPUMP" +} +account.login(**params) + +account.signal.connect() +with account.signal.expect("messages.new") as exp: + account.send_message(chat_id, "hello") + +# Available only after the `with` block exits +print(exp.result) + +account.signal.disconnect() +``` + +**Note**: Some signals arrive faster than others - the Status Backend does not emit every signal type at the same speed (for example, a local `envelope.sent` confirmation typically arrives well before a `messages.new` event that depends on network propagation). Tune the `timeout` per signal type, and use `count` when an action is expected to produce multiple signals, rather than assuming they all land at once or in a fixed order. + ### `logger`