diff --git a/aikido_zen/background_process/cloud_connection_manager/__init__.py b/aikido_zen/background_process/cloud_connection_manager/__init__.py index ea92f5074..4e69037c4 100644 --- a/aikido_zen/background_process/cloud_connection_manager/__init__.py +++ b/aikido_zen/background_process/cloud_connection_manager/__init__.py @@ -1,4 +1,4 @@ -""" This file simply exports the CloudConnectionManager class""" +"""This file simply exports the CloudConnectionManager class""" from aikido_zen.background_process.heartbeats import send_heartbeats_every_x_secs from aikido_zen.background_process.routes import Routes @@ -10,10 +10,12 @@ from aikido_zen.storage.users import Users from aikido_zen.storage.hostnames import Hostnames from ..realtime.start_polling_for_changes import start_polling_for_changes +from ..realtime.listen_for_config_updates import listen_for_config_updates from ...helpers.get_current_unixtime_ms import get_unixtime_ms from ...storage.ai_statistics import AIStatistics from ...storage.firewall_lists import FirewallLists from ...storage.statistics import Statistics +from aikido_zen.helpers.env_vars.feature_flags import is_feature_enabled # Import functions : from .get_manager_info import get_manager_info @@ -67,6 +69,9 @@ def start(self, event_scheduler): send_heartbeats_every_x_secs(self, self.heartbeat_secs, event_scheduler) start_polling_for_changes(self, event_scheduler) + if is_feature_enabled("sse"): + listen_for_config_updates(self, event_scheduler) + def report_initial_stats(self): """ This is run 1m after startup, and checks if we should send out diff --git a/aikido_zen/background_process/realtime/__init__.py b/aikido_zen/background_process/realtime/__init__.py index e50e08106..5f12f58bf 100644 --- a/aikido_zen/background_process/realtime/__init__.py +++ b/aikido_zen/background_process/realtime/__init__.py @@ -15,7 +15,7 @@ def get_realtime_url(): if not realtime_url.endswith("/"): realtime_url += "/" return realtime_url - return "https://runtime.aikido.dev/" + return get_api_url() def get_config(token): diff --git a/aikido_zen/background_process/realtime/listen_for_config_updates.py b/aikido_zen/background_process/realtime/listen_for_config_updates.py new file mode 100644 index 000000000..5253a4d18 --- /dev/null +++ b/aikido_zen/background_process/realtime/listen_for_config_updates.py @@ -0,0 +1,59 @@ +""" +Mainly exports `listen_for_config_updates` +""" + +import json + +from aikido_zen.helpers.token import Token +from aikido_zen.helpers.logging import logger +import aikido_zen.background_process.realtime as realtime +from .sse_client import connect_to_sse + + +def listen_for_config_updates(connection_manager, event_scheduler): + """ + Connects to the realtime SSE endpoint and fetches the new config whenever + the server signals, through a "config-updated" event, that it has changed. + """ + if not isinstance(connection_manager.token, Token): + logger.info("No token provided, not listening for config updates") + return + if connection_manager.serverless: + logger.info( + "Running in serverless environment, not listening for config updates" + ) + return + + token = connection_manager.token + last_updated_at = connection_manager.conf.last_updated_at + + def on_event(event): + nonlocal last_updated_at + + logger.debug("SSE event received: %s", event.event) + if event.event != "config-updated": + return + + try: + payload = json.loads(event.data) + config_updated_at = payload["configUpdatedAt"] + if config_updated_at <= last_updated_at: + return + except (ValueError, KeyError, TypeError): + logger.debug("SSE config-updated event has invalid payload: %s", event.data) + return + + logger.debug("SSE config-updated event, fetching new config") + + try: + config = realtime.get_config(token) + logger.debug( + "SSE config fetched, configUpdatedAt: %s", config.get("configUpdatedAt") + ) + last_updated_at = config.get("configUpdatedAt", config_updated_at) + connection_manager.update_service_config({**config, "success": True}) + connection_manager.update_firewall_lists() + except Exception as e: + logger.error("Failed to fetch config after SSE event : %s", e) + + connect_to_sse(token=token, on_event=on_event) diff --git a/aikido_zen/background_process/realtime/listen_for_config_updates_test.py b/aikido_zen/background_process/realtime/listen_for_config_updates_test.py new file mode 100644 index 000000000..575f6d7c6 --- /dev/null +++ b/aikido_zen/background_process/realtime/listen_for_config_updates_test.py @@ -0,0 +1,169 @@ +import json +import logging +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from aikido_zen.helpers.token import Token +from .listen_for_config_updates import listen_for_config_updates + + +def make_event(event="config-updated", data=None): + return SimpleNamespace( + event=event, data=json.dumps(data) if data is not None else "" + ) + + +@pytest.fixture +def connection_manager(): + return MagicMock( + token=Token("123"), + serverless=None, + conf=MagicMock(last_updated_at=0), + ) + + +def test_no_token(caplog): + with patch( + "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse" + ) as mock_connect: + listen_for_config_updates( + connection_manager=MagicMock(token=None, serverless=None), + event_scheduler=MagicMock(), + ) + + assert "No token provided, not listening for config updates" in caplog.text + mock_connect.assert_not_called() + + +def test_serverless_environment(caplog): + with patch( + "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse" + ) as mock_connect: + listen_for_config_updates( + connection_manager=MagicMock(token=Token("123"), serverless=True), + event_scheduler=MagicMock(), + ) + + assert ( + "Running in serverless environment, not listening for config updates" + in caplog.text + ) + mock_connect.assert_not_called() + + +def test_connects_to_sse_with_token(connection_manager): + with patch( + "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse" + ) as mock_connect: + listen_for_config_updates( + connection_manager=connection_manager, event_scheduler=MagicMock() + ) + + mock_connect.assert_called_once() + assert mock_connect.call_args.kwargs["token"] is connection_manager.token + assert callable(mock_connect.call_args.kwargs["on_event"]) + + +def get_on_event(connection_manager): + with patch( + "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse" + ) as mock_connect: + listen_for_config_updates( + connection_manager=connection_manager, event_scheduler=MagicMock() + ) + return mock_connect.call_args.kwargs["on_event"] + + +def test_ignores_events_that_are_not_config_updated(connection_manager): + on_event = get_on_event(connection_manager) + + with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config: + on_event(make_event(event="ping")) + + mock_get_config.assert_not_called() + connection_manager.update_service_config.assert_not_called() + + +def test_ignores_config_updated_event_with_invalid_json(connection_manager, caplog): + caplog.set_level(logging.DEBUG, logger="Zen") + on_event = get_on_event(connection_manager) + + with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config: + on_event(SimpleNamespace(event="config-updated", data="not json")) + + mock_get_config.assert_not_called() + connection_manager.update_service_config.assert_not_called() + assert "SSE config-updated event has invalid payload" in caplog.text + + +def test_ignores_config_updated_event_missing_config_updated_at(connection_manager): + on_event = get_on_event(connection_manager) + + with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config: + on_event(make_event(data={"foo": "bar"})) + + mock_get_config.assert_not_called() + connection_manager.update_service_config.assert_not_called() + + +def test_ignores_config_updated_event_that_is_not_newer(connection_manager): + connection_manager.conf.last_updated_at = 100 + on_event = get_on_event(connection_manager) + + with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config: + on_event(make_event(data={"configUpdatedAt": 100})) + + mock_get_config.assert_not_called() + connection_manager.update_service_config.assert_not_called() + + +def test_fetches_and_applies_new_config_on_newer_event(connection_manager): + connection_manager.conf.last_updated_at = 100 + on_event = get_on_event(connection_manager) + + new_config = {"endpoints": [], "configUpdatedAt": 200} + with patch( + "aikido_zen.background_process.realtime.get_config", return_value=new_config + ) as mock_get_config: + on_event(make_event(data={"configUpdatedAt": 200})) + + mock_get_config.assert_called_once_with(connection_manager.token) + connection_manager.update_service_config.assert_called_once_with( + {**new_config, "success": True} + ) + connection_manager.update_firewall_lists.assert_called_once() + + +def test_updates_last_updated_at_so_a_second_stale_event_is_ignored( + connection_manager, +): + connection_manager.conf.last_updated_at = 100 + on_event = get_on_event(connection_manager) + + new_config = {"endpoints": [], "configUpdatedAt": 200} + with patch( + "aikido_zen.background_process.realtime.get_config", return_value=new_config + ) as mock_get_config: + on_event(make_event(data={"configUpdatedAt": 200})) + # A second event claiming to be newer than the original lastUpdatedAt + # (100), but not newer than what we just fetched (200), should now be + # ignored without fetching again. + on_event(make_event(data={"configUpdatedAt": 150})) + + mock_get_config.assert_called_once() + + +def test_handles_get_config_failure_gracefully(connection_manager, caplog): + connection_manager.conf.last_updated_at = 100 + on_event = get_on_event(connection_manager) + + with patch( + "aikido_zen.background_process.realtime.get_config", + side_effect=ValueError("Request timed out"), + ): + on_event(make_event(data={"configUpdatedAt": 200})) + + connection_manager.update_service_config.assert_not_called() + assert "Failed to fetch config after SSE event" in caplog.text diff --git a/aikido_zen/background_process/realtime/sse_client/__init__.py b/aikido_zen/background_process/realtime/sse_client/__init__.py new file mode 100644 index 000000000..d46d3bd81 --- /dev/null +++ b/aikido_zen/background_process/realtime/sse_client/__init__.py @@ -0,0 +1,130 @@ +""" +SSE (Server-Sent Events) client that connects to the realtime endpoint and +dispatches events, automatically reconnecting with exponential backoff and +jitter on disconnects. +""" + +import http.client +import random +import socket +import threading +import time +import urllib.error +import urllib.request + +import aikido_zen.background_process.realtime as realtime +from aikido_zen.helpers.logging import logger +from .parser import SSEParser +import aikido_zen.config as config + +INITIAL_RECONNECT_SECS = 5 +MAX_RECONNECT_SECS = 60 +STABLE_CONNECTION_SECS = 30 +READ_TIMEOUT_SECS = 70 + +_CONNECTION_ERRORS = (urllib.error.URLError, OSError, http.client.HTTPException) + + +def connect_to_sse( + token, + on_event, + initial_reconnect_secs=INITIAL_RECONNECT_SECS, + read_timeout_secs=READ_TIMEOUT_SECS, +): + """ + Starts a daemon thread that connects to the realtime SSE endpoint and calls + `on_event(event)` for every event received. Reconnects with exponential + backoff (and jitter) on disconnect, until the server rejects the connection + with a 401 or 403 status, at which point it stops. + + Returns the thread that was started. + """ + thread = threading.Thread( + target=_reconnect_loop, + args=(token, on_event, initial_reconnect_secs, read_timeout_secs), + daemon=True, + ) + thread.start() + return thread + + +def _reconnect_loop(token, on_event, initial_reconnect_secs, read_timeout_secs): + reconnect_secs = initial_reconnect_secs + try: + while True: + start = time.monotonic() + outcome, status_code = _connect(token, on_event, read_timeout_secs) + + if outcome == "disconnected" and status_code in (401, 403): + logger.info( + "SSE connection rejected with status %s, stopping", status_code + ) + return + + if time.monotonic() - start >= STABLE_CONNECTION_SECS: + reconnect_secs = initial_reconnect_secs + + jitter = random.random() * (reconnect_secs / 2) + delay_secs = reconnect_secs + jitter + + logger.debug("SSE scheduling reconnect in %sms", round(delay_secs * 1000)) + + reconnect_secs = min(reconnect_secs * 2, MAX_RECONNECT_SECS) + + time.sleep(delay_secs) + except Exception as e: + logger.error("SSE loop error : %s", e) + + +def _connect(token, on_event, read_timeout_secs): + """ + Opens a single SSE connection and dispatches events until the connection is + closed, errors out, or goes idle for longer than `read_timeout_secs`. + + Returns a tuple of (outcome, status_code), where outcome is either "error" + (the connection could not be made, or timed out) or "disconnected" (a + response was received, but the connection has since ended). + """ + url = f"{realtime.get_realtime_url()}api/runtime/stream" + logger.debug("SSE connecting to %s", url) + + request = urllib.request.Request( + url, + method="GET", + headers={ + "Authorization": str(token), + "Accept": "text/event-stream", + "Cache-Control": "no-cache", + "X-Agent-Platform": "python", + "X-Agent-Version": config.PKG_VERSION, + }, + ) + + try: + with urllib.request.urlopen(request, timeout=read_timeout_secs) as response: + status_code = response.getcode() + if status_code != 200: + return "disconnected", status_code + + logger.debug("SSE connected successfully") + + parser = SSEParser(response) + try: + for event in parser.events(): + logger.debug("SSE received event : %s", event) + on_event(event) + except socket.timeout: + logger.debug("SSE read timeout") + return "error", None + except _CONNECTION_ERRORS as e: + logger.debug("SSE stream error : %s", e) + return "disconnected", status_code + except urllib.error.HTTPError as e: + logger.debug("SSE connection rejected with status %s", e.code) + return "disconnected", e.code + except _CONNECTION_ERRORS as e: + logger.debug("SSE connection error : %s", e) + return "error", None + + logger.debug("SSE connection closed by server") + return "disconnected", status_code diff --git a/aikido_zen/background_process/realtime/sse_client/init_test.py b/aikido_zen/background_process/realtime/sse_client/init_test.py new file mode 100644 index 000000000..9e72954c6 --- /dev/null +++ b/aikido_zen/background_process/realtime/sse_client/init_test.py @@ -0,0 +1,221 @@ +import socket +import urllib.error +from unittest.mock import MagicMock, patch + +from aikido_zen.helpers.token import Token +from . import ( + _connect, + _reconnect_loop, + connect_to_sse, + INITIAL_RECONNECT_SECS, + READ_TIMEOUT_SECS, +) + +REALTIME_URL_PATCH = patch( + "aikido_zen.background_process.realtime.get_realtime_url", + return_value="https://runtime.aikido.dev/", +) + + +def make_response(status_code, chunks=None): + response = MagicMock() + response.getcode.return_value = status_code + response.__enter__.return_value = response + response.__exit__.side_effect = lambda *args: response.close() + if chunks is not None: + response.__iter__.return_value = iter(chunks) + return response + + +def test_connect_dispatches_events_and_reports_disconnected_on_stream_end(): + response = make_response(200, [b"data: hello\n\n", b"data: world\n\n"]) + events = [] + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", return_value=response): + outcome, status_code = _connect(Token("123"), events.append, 5) + + assert outcome == "disconnected" + assert status_code == 200 + assert [e.data for e in events] == ["hello", "world"] + response.close.assert_called_once() # closed via parser.close() + + +def test_connect_sends_expected_request(): + response = make_response(200, []) + captured_request = {} + + def fake_urlopen(request, timeout): + captured_request["request"] = request + captured_request["timeout"] = timeout + return response + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", side_effect=fake_urlopen): + _connect(Token("my-token"), MagicMock(), 42) + + request = captured_request["request"] + assert request.full_url == "https://runtime.aikido.dev/api/runtime/stream" + assert request.get_header("Authorization") == "my-token" + assert request.get_header("Accept") == "text/event-stream" + assert captured_request["timeout"] == 42 + + +def test_connect_non_200_status_closes_response_without_reading_body(): + response = make_response(500) + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", return_value=response): + outcome, status_code = _connect(Token("123"), MagicMock(), 5) + + assert outcome == "disconnected" + assert status_code == 500 + response.close.assert_called_once() + + +def test_connect_http_error_returns_disconnected_with_status(): + http_error = urllib.error.HTTPError( + url="https://runtime.aikido.dev/api/runtime/stream", + code=403, + msg="Forbidden", + hdrs=None, + fp=None, + ) + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", side_effect=http_error): + outcome, status_code = _connect(Token("123"), MagicMock(), 5) + + assert outcome == "disconnected" + assert status_code == 403 + + +def test_connect_connection_error_returns_error_outcome(): + with REALTIME_URL_PATCH, patch( + "urllib.request.urlopen", + side_effect=urllib.error.URLError("connection refused"), + ): + outcome, status_code = _connect(Token("123"), MagicMock(), 5) + + assert outcome == "error" + assert status_code is None + + +def test_connect_read_timeout_mid_stream_returns_error(): + response = make_response(200) + response.__iter__.side_effect = socket.timeout("timed out") + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", return_value=response): + outcome, status_code = _connect(Token("123"), MagicMock(), 5) + + assert outcome == "error" + assert status_code is None + response.close.assert_called_once() + + +def test_connect_generic_stream_error_returns_disconnected_with_status(): + response = make_response(200) + response.__iter__.side_effect = ConnectionResetError("reset") + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", return_value=response): + outcome, status_code = _connect(Token("123"), MagicMock(), 5) + + assert outcome == "disconnected" + assert status_code == 200 + + +def test_reconnect_loop_stops_immediately_on_401_or_403(): + for status_code in (401, 403): + connect_calls = [] + + def fake_connect(token, on_event, read_timeout_secs): + connect_calls.append(status_code) + return "disconnected", status_code + + with patch( + "aikido_zen.background_process.realtime.sse_client._connect", + side_effect=fake_connect, + ), patch( + "aikido_zen.background_process.realtime.sse_client.time.sleep" + ) as mock_sleep: + _reconnect_loop(Token("123"), MagicMock(), INITIAL_RECONNECT_SECS, 5) + + assert connect_calls == [status_code] + mock_sleep.assert_not_called() + + +def test_reconnect_loop_backoff_doubles_between_reconnects(): + outcomes = iter( + [("error", None), ("error", None), ("disconnected", 401)], + ) + sleep_calls = [] + monotonic_values = iter([0, 0, 10, 10, 20]) + + with patch( + "aikido_zen.background_process.realtime.sse_client._connect", + side_effect=lambda *a: next(outcomes), + ), patch( + "aikido_zen.background_process.realtime.sse_client.time.sleep", + side_effect=sleep_calls.append, + ), patch( + "aikido_zen.background_process.realtime.sse_client.time.monotonic", + side_effect=lambda: next(monotonic_values), + ), patch( + "aikido_zen.background_process.realtime.sse_client.random.random", + return_value=0, + ): + _reconnect_loop(Token("123"), MagicMock(), 5, 5) + + assert sleep_calls == [5, 10] + + +def test_reconnect_loop_resets_backoff_after_stable_connection(): + outcomes = iter( + [ + ("error", None), + ("error", None), + ("error", None), + ("disconnected", 401), + ], + ) + sleep_calls = [] + # start1=0 end1=0 (unstable) | start2=10 end2=10 (unstable) | start3=20 end3=60 (stable, resets) | start4=unused + monotonic_values = iter([0, 0, 10, 10, 20, 60, 999]) + + with patch( + "aikido_zen.background_process.realtime.sse_client._connect", + side_effect=lambda *a: next(outcomes), + ), patch( + "aikido_zen.background_process.realtime.sse_client.time.sleep", + side_effect=sleep_calls.append, + ), patch( + "aikido_zen.background_process.realtime.sse_client.time.monotonic", + side_effect=lambda: next(monotonic_values), + ), patch( + "aikido_zen.background_process.realtime.sse_client.random.random", + return_value=0, + ): + _reconnect_loop(Token("123"), MagicMock(), 5, 5) + + assert sleep_calls == [5, 10, 5] + + +def test_reconnect_loop_logs_unexpected_exception(caplog): + with patch( + "aikido_zen.background_process.realtime.sse_client._connect", + side_effect=RuntimeError("boom"), + ): + _reconnect_loop(Token("123"), MagicMock(), 5, 5) + + assert "SSE loop error : boom" in caplog.text + + +def test_connect_to_sse_starts_a_daemon_thread_running_the_reconnect_loop(): + with patch( + "aikido_zen.background_process.realtime.sse_client._reconnect_loop" + ) as mock_loop: + token = Token("123") + on_event = MagicMock() + thread = connect_to_sse(token, on_event) + thread.join(timeout=2) + + assert thread.daemon is True + mock_loop.assert_called_once_with( + token, on_event, INITIAL_RECONNECT_SECS, READ_TIMEOUT_SECS + ) diff --git a/aikido_zen/background_process/realtime/sse_client/parser.py b/aikido_zen/background_process/realtime/sse_client/parser.py new file mode 100644 index 000000000..c6e48c607 --- /dev/null +++ b/aikido_zen/background_process/realtime/sse_client/parser.py @@ -0,0 +1,129 @@ +""" +Based on https://github.com/mpetazzoni/sseclient +Licensed under the Apache License, Version 2.0 +Copyright Maxime Petazzoni + +Server Side Events (SSE) client for Python. + +Provides a generator of SSE received through an existing HTTP response. +""" + +_FIELD_SEPARATOR = ":" + + +class SSEParser: + """Implementation of a SSE parser. + + See http://www.w3.org/TR/2009/WD-eventsource-20091029/ for the + specification. + """ + + def __init__(self, event_source, char_enc="utf-8"): + """Initialize the SSE parser over an existing, ready to consume + event source. + + The event source is expected to be a binary stream and have a close() + method. That would usually be something that implements + io.BinaryIOBase, like an httplib or urllib3 HTTPResponse object. + """ + self._event_source = event_source + self._char_enc = char_enc + + def _read(self): + """Read the incoming event source stream and yield event chunks. + + Unfortunately it is possible for some servers to decide to break an + event into multiple HTTP chunks in the response. It is thus necessary + to correctly stitch together consecutive response chunks and find the + SSE delimiter (empty new line) to yield full, correct event chunks.""" + buffer = [] + tail = b"" + for chunk in self._event_source: + for line in chunk.splitlines(True): + buffer.append(line) + tail = (tail + line)[-4:] + if tail.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")): + yield b"".join(buffer) + buffer = [] + tail = b"" + if buffer: + yield b"".join(buffer) + + def events(self): + for chunk in self._read(): + event = Event() + # Split before decoding so splitlines() only uses \r and \n + for line in chunk.splitlines(): + # Decode the line. + line = line.decode(self._char_enc) + + # SSE spec: skip empty lines and colon-prefixed comment lines. + if not line.strip() or line.startswith(_FIELD_SEPARATOR): + continue + + data = line.split(_FIELD_SEPARATOR, 1) + field = data[0] + + # Ignore unknown fields. + if field not in event.__dict__: + continue + + if len(data) > 1: + # From the spec: + # "If value starts with a single U+0020 SPACE character, + # remove it from value." + if data[1].startswith(" "): + value = data[1][1:] + else: + value = data[1] + else: + # If no value is present after the separator, + # assume an empty value. + value = "" + + # The data field may come over multiple lines and their values + # are concatenated with each other. + if field == "data": + event.__dict__[field] += value + "\n" + else: + event.__dict__[field] = value + + # Events with no data are not dispatched. + if not event.data: + continue + + # If the data field ends with a newline, remove it. + if event.data.endswith("\n"): + event.data = event.data[0:-1] + + # Empty event names default to 'message' + event.event = event.event or "message" + + # Dispatch the event + yield event + + def close(self): + """Manually close the event source stream.""" + self._event_source.close() + + +class Event: + """Representation of an event from the event stream.""" + + def __init__(self, event_id=None, event="message", data="", retry=None): + self.id = event_id + self.event = event + self.data = data + self.retry = retry + + def __str__(self): + s = f"{self.event} event" + if self.id: + s += f" #{self.id}" + if self.data: + s += f", {len(self.data)} byte{'s' if len(self.data) != 1 else ''}" + else: + s += ", no data" + if self.retry: + s += f", retry in {self.retry}ms" + return s diff --git a/aikido_zen/background_process/realtime/sse_client/parser_test.py b/aikido_zen/background_process/realtime/sse_client/parser_test.py new file mode 100644 index 000000000..e9765c93e --- /dev/null +++ b/aikido_zen/background_process/realtime/sse_client/parser_test.py @@ -0,0 +1,186 @@ +from unittest.mock import MagicMock +from .parser import SSEParser, Event + + +def make_source(chunks): + """Creates an iterable of byte chunks with a close() method, like an HTTP response.""" + source = MagicMock() + source.__iter__.return_value = iter(chunks) + return source + + +def test_read_single_chunk_single_event(): + source = make_source([b"data: hello\n\n"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: hello\n\n"] + + +def test_read_splits_multiple_events_in_one_chunk(): + source = make_source([b"data: one\n\ndata: two\n\n"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: one\n\n", b"data: two\n\n"] + + +def test_read_stitches_event_split_across_chunks(): + source = make_source([b"data: hel", b"lo\n\n"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: hello\n\n"] + + +def test_read_handles_carriage_return_delimiters(): + source = make_source([b"data: hello\r\r"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: hello\r\r"] + + +def test_read_handles_crlf_delimiters(): + source = make_source([b"data: hello\r\n\r\n"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: hello\r\n\r\n"] + + +def test_read_yields_trailing_incomplete_chunk(): + source = make_source([b"data: hello"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: hello"] + + +def test_events_basic_data_field(): + source = make_source([b"data: hello\n\n"]) + parser = SSEParser(source) + + events = list(parser.events()) + + assert len(events) == 1 + assert events[0].data == "hello" + assert events[0].event == "message" + + +def test_events_defaults_event_name_to_message(): + source = make_source([b"data: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.event == "message" + + +def test_events_custom_event_name_and_id(): + source = make_source([b"event: update\nid: 42\ndata: payload\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.event == "update" + assert event.id == "42" + assert event.data == "payload" + + +def test_events_multiline_data_is_joined_with_newlines(): + source = make_source([b"data: line1\ndata: line2\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.data == "line1\nline2" + + +def test_events_ignores_comment_lines(): + source = make_source([b": this is a comment\ndata: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.data == "hello" + + +def test_events_ignores_unknown_fields(): + source = make_source([b"foo: bar\ndata: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.data == "hello" + assert not hasattr(event, "foo") + + +def test_events_strips_single_leading_space_from_value(): + source = make_source([b"data: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.data == "hello" + + +def test_events_keeps_extra_leading_spaces_beyond_first(): + source = make_source([b"data: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.data == " hello" + + +def test_events_field_with_no_value_is_empty_string(): + source = make_source([b"event\ndata: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.event == "message" + + +def test_events_skips_events_with_no_data(): + source = make_source([b"event: ping\n\ndata: hello\n\n"]) + parser = SSEParser(source) + + events = list(parser.events()) + assert len(events) == 1 + assert events[0].data == "hello" + + +def test_events_multiple_events_in_stream(): + source = make_source([b"data: first\n\ndata: second\n\n"]) + parser = SSEParser(source) + + events = list(parser.events()) + assert [e.data for e in events] == ["first", "second"] + + +def test_events_retry_field(): + source = make_source([b"retry: 5000\ndata: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.retry == "5000" + + +def test_events_uses_configured_char_encoding(): + source = make_source(["data: héllo\n\n".encode("latin-1")]) + parser = SSEParser(source, char_enc="latin-1") + + [event] = parser.events() + assert event.data == "héllo" + + +def test_close_delegates_to_event_source(): + source = make_source([]) + parser = SSEParser(source) + + parser.close() + + source.close.assert_called_once() + + +def test_event_str_with_data_and_id(): + event = Event(event_id="1", event="message", data="hello") + assert str(event) == "message event #1, 5 bytes" + + +def test_event_str_without_data(): + event = Event() + assert str(event) == "message event, no data" + + +def test_event_str_with_retry(): + event = Event(data="hi", retry=3000) + assert str(event) == "message event, 2 bytes, retry in 3000ms" diff --git a/aikido_zen/background_process/realtime/sse_client/real_server_test.py b/aikido_zen/background_process/realtime/sse_client/real_server_test.py new file mode 100644 index 000000000..7a216879a --- /dev/null +++ b/aikido_zen/background_process/realtime/sse_client/real_server_test.py @@ -0,0 +1,224 @@ +""" +Integration tests that exercise the SSE client against a real HTTP server +(as opposed to init_test.py, which mocks urllib.request.urlopen), so that real +socket/timeout/connection-refused behavior is covered, not just our own mocks. +""" + +import http.server +import json +import threading +import time +from unittest.mock import patch + +from aikido_zen.helpers.token import Token +from . import _connect, _reconnect_loop + +REALTIME_URL_PATCH_TARGET = "aikido_zen.background_process.realtime.get_realtime_url" + + +class _TestServer(http.server.ThreadingHTTPServer): + daemon_threads = True # Don't let lingering handler threads block process exit + + +def make_handler(do_get): + """Builds a BaseHTTPRequestHandler class that delegates GET handling to `do_get(handler)`.""" + + class Handler(http.server.BaseHTTPRequestHandler): + protocol_version = "HTTP/1.0" + + def do_GET(self): # pylint: disable=invalid-name + do_get(self) + + def log_message(self, *args, **kwargs): + pass + + return Handler + + +def start_server(handler_class): + server = _TestServer(("127.0.0.1", 0), handler_class) + threading.Thread(target=server.serve_forever, daemon=True).start() + return server, f"http://127.0.0.1:{server.server_address[1]}/" + + +def stop_server(server): + server.shutdown() + server.server_close() + + +def wait_until(predicate, timeout=2.0, interval=0.01): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(interval) + return predicate() + + +def test_happy_path_receives_events_with_auth_header(): + captured = {} + events = [] + + def do_get(handler): + captured["auth"] = handler.headers.get("Authorization") + captured["accept"] = handler.headers.get("Accept") + captured["cache"] = handler.headers.get("Cache-Control") + captured["x-agent-platform"] = handler.headers.get("X-Agent-Platform") + captured["x-agent-version"] = handler.headers.get("X-Agent-Version") + handler.send_response(200) + handler.send_header("Content-Type", "text/event-stream") + handler.send_header("Cache-Control", "no-cache") + handler.end_headers() + handler.wfile.write(b": ping\n\n") + handler.wfile.write( + b'event: config-updated\ndata: {"configUpdatedAt": 100}\n\n' + ) + handler.wfile.flush() + time.sleep(0.2) + handler.wfile.write( + b'event: config-updated\ndata: {"configUpdatedAt": 200}\n\n' + ) + handler.wfile.flush() + + server, url = start_server(make_handler(do_get)) + try: + with patch(REALTIME_URL_PATCH_TARGET, return_value=url): + outcome, status_code = _connect( + Token("my-secret-token"), events.append, read_timeout_secs=5 + ) + finally: + stop_server(server) + + assert outcome == "disconnected" + assert status_code == 200 + assert captured["auth"] == "my-secret-token" + assert captured["accept"] == "text/event-stream" + assert captured["cache"] == "no-cache" + assert captured["x-agent-platform"] == "python" + assert captured["x-agent-version"] == "1.0-REPLACE-VERSION" + assert len(events) == 2 + assert events[0].event == "config-updated" + assert json.loads(events[0].data) == {"configUpdatedAt": 100} + assert json.loads(events[1].data) == {"configUpdatedAt": 200} + + +def test_stops_reconnecting_on_401(): + connection_count = {"n": 0} + + def do_get(handler): + connection_count["n"] += 1 + handler.send_response(401) + handler.end_headers() + + server, url = start_server(make_handler(do_get)) + try: + with patch(REALTIME_URL_PATCH_TARGET, return_value=url): + _reconnect_loop( + Token("bad-token"), + lambda event: None, + initial_reconnect_secs=0.05, + read_timeout_secs=1, + ) + finally: + stop_server(server) + + assert connection_count["n"] == 1 + + +def test_reconnects_after_non_200_status(): + connection_count = {"n": 0} + + def do_get(handler): + connection_count["n"] += 1 + if connection_count["n"] == 1: + handler.send_response(500) + handler.end_headers() + return + handler.send_response(200) + handler.send_header("Content-Type", "text/event-stream") + handler.send_header("Cache-Control", "no-cache") + handler.end_headers() + handler.wfile.write(b": ping\n\n") + handler.wfile.flush() + + server, url = start_server(make_handler(do_get)) + try: + with patch(REALTIME_URL_PATCH_TARGET, return_value=url): + thread = threading.Thread( + target=_reconnect_loop, + args=(Token("test-token"), lambda event: None, 0.05, 1), + daemon=True, + ) + thread.start() + assert wait_until(lambda: connection_count["n"] >= 2) + finally: + stop_server(server) + + +def test_connection_refused(): + probe = _TestServer(("127.0.0.1", 0), make_handler(lambda handler: None)) + port = probe.server_address[1] + probe.server_close() + + with patch(REALTIME_URL_PATCH_TARGET, return_value=f"http://127.0.0.1:{port}/"): + outcome, status_code = _connect( + Token("test-token"), lambda event: None, read_timeout_secs=1 + ) + + assert outcome == "error" + assert status_code is None + + +def test_reconnects_on_read_timeout(): + connection_count = {"n": 0} + + def do_get(handler): + connection_count["n"] += 1 + handler.send_response(200) + handler.send_header("Content-Type", "text/event-stream") + handler.send_header("Cache-Control", "no-cache") + handler.end_headers() + if connection_count["n"] == 1: + time.sleep(1) + else: + handler.wfile.write(b": ping\n\n") + handler.wfile.flush() + + server, url = start_server(make_handler(do_get)) + try: + with patch(REALTIME_URL_PATCH_TARGET, return_value=url): + thread = threading.Thread( + target=_reconnect_loop, + args=(Token("test-token"), lambda event: None, 0.05, 0.2), + daemon=True, + ) + thread.start() + assert wait_until(lambda: connection_count["n"] >= 2) + finally: + stop_server(server) + + +def test_reconnects_when_server_closes_connection(): + connection_count = {"n": 0} + + def do_get(handler): + connection_count["n"] += 1 + handler.send_response(200) + handler.send_header("Content-Type", "text/event-stream") + handler.send_header("Cache-Control", "no-cache") + handler.end_headers() + handler.wfile.write(b": ping\n\n") + handler.wfile.flush() + + server, url = start_server(make_handler(do_get)) + try: + with patch(REALTIME_URL_PATCH_TARGET, return_value=url): + thread = threading.Thread( + target=_reconnect_loop, + args=(Token("test-token"), lambda event: None, 0.05, 1), + daemon=True, + ) + thread.start() + assert wait_until(lambda: connection_count["n"] >= 2) + finally: + stop_server(server)