From 902a8384a35fa03c654e4b9250865e243328c44b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6ssler?= Date: Thu, 30 Jul 2026 15:21:14 +0200 Subject: [PATCH 1/2] Add support for enabled features in config --- .../cloud_connection_manager/__init__.py | 4 +- .../cloud_connection_manager/init_test.py | 54 +++++++++++++++++++ .../update_service_config.py | 3 ++ .../update_service_config_test.py | 45 ++++++++++++++++ .../background_process/service_config.py | 9 ++++ .../background_process/service_config_test.py | 22 ++++++++ 6 files changed, 136 insertions(+), 1 deletion(-) diff --git a/aikido_zen/background_process/cloud_connection_manager/__init__.py b/aikido_zen/background_process/cloud_connection_manager/__init__.py index 4e69037c4..658b6e76b 100644 --- a/aikido_zen/background_process/cloud_connection_manager/__init__.py +++ b/aikido_zen/background_process/cloud_connection_manager/__init__.py @@ -69,7 +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"): + if is_feature_enabled("sse") or self.conf.is_feature_enabled( + "realtime_updates" + ): listen_for_config_updates(self, event_scheduler) def report_initial_stats(self): diff --git a/aikido_zen/background_process/cloud_connection_manager/init_test.py b/aikido_zen/background_process/cloud_connection_manager/init_test.py index 3c58a85af..e14751904 100644 --- a/aikido_zen/background_process/cloud_connection_manager/init_test.py +++ b/aikido_zen/background_process/cloud_connection_manager/init_test.py @@ -1,4 +1,5 @@ import pytest +from unittest.mock import patch, MagicMock from aikido_zen.helpers.token import Token from aikido_zen.background_process.api.http_api import ReportingApiHTTP from aikido_zen.background_process.service_config import ServiceConfig @@ -45,4 +46,57 @@ def test_cloud_connection_manager_empty_serverless(): CloudConnectionManager(block, api, token, serverless) +@patch("aikido_zen.background_process.cloud_connection_manager.on_start") +@patch( + "aikido_zen.background_process.cloud_connection_manager.start_polling_for_changes" +) +@patch( + "aikido_zen.background_process.cloud_connection_manager.send_heartbeats_every_x_secs" +) +@patch( + "aikido_zen.background_process.cloud_connection_manager.listen_for_config_updates" +) +def test_start_enables_sse_when_server_flag_set( + mock_listen_for_config_updates, + mock_send_heartbeats, + mock_start_polling, + mock_on_start, + setup_cloud_connection_manager, +): + """SSE listening should start when the server enables the realtime_updates feature flag""" + manager = setup_cloud_connection_manager + mock_on_start.return_value = {"success": True} + manager.conf.update_enabled_features(["realtime_updates"]) + + manager.start(event_scheduler=MagicMock()) + + mock_listen_for_config_updates.assert_called_once() + + +@patch("aikido_zen.background_process.cloud_connection_manager.on_start") +@patch( + "aikido_zen.background_process.cloud_connection_manager.start_polling_for_changes" +) +@patch( + "aikido_zen.background_process.cloud_connection_manager.send_heartbeats_every_x_secs" +) +@patch( + "aikido_zen.background_process.cloud_connection_manager.listen_for_config_updates" +) +def test_start_does_not_enable_sse_without_flag( + mock_listen_for_config_updates, + mock_send_heartbeats, + mock_start_polling, + mock_on_start, + setup_cloud_connection_manager, +): + """SSE listening should not start when neither the env var nor the server flag is set""" + manager = setup_cloud_connection_manager + mock_on_start.return_value = {"success": True} + + manager.start(event_scheduler=MagicMock()) + + mock_listen_for_config_updates.assert_not_called() + + # Additional tests can be added here for other edge cases or scenarios diff --git a/aikido_zen/background_process/cloud_connection_manager/update_service_config.py b/aikido_zen/background_process/cloud_connection_manager/update_service_config.py index 4184d42b6..e3b351148 100644 --- a/aikido_zen/background_process/cloud_connection_manager/update_service_config.py +++ b/aikido_zen/background_process/cloud_connection_manager/update_service_config.py @@ -38,3 +38,6 @@ def update_service_config(connection_manager, res): if "domains" in res: connection_manager.conf.update_outbound_domains(res["domains"]) + + if isinstance(res.get("enabledFeatures"), list): + connection_manager.conf.update_enabled_features(res["enabledFeatures"]) diff --git a/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py b/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py index 6080a3515..e7c815bfa 100644 --- a/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py +++ b/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py @@ -260,6 +260,51 @@ def test_update_service_config_excluded_user_ids_from_rate_limiting(): assert connection_manager.conf.is_user_excluded_from_rate_limiting("user3") is False +def test_update_service_config_enabled_features(): + """Test that update_service_config handles enabledFeatures""" + connection_manager = MagicMock() + connection_manager.conf = ServiceConfig( + endpoints=[], + last_updated_at=0, + blocked_uids=set(), + bypassed_ips=[], + received_any_stats=False, + ) + connection_manager.block = False + + res = { + "success": True, + "enabledFeatures": ["realtime_updates"], + } + + update_service_config(connection_manager, res) + + assert connection_manager.conf.is_feature_enabled("realtime_updates") is True + assert connection_manager.conf.is_feature_enabled("some_other_feature") is False + + +def test_update_service_config_enabled_features_not_array(): + """Test that update_service_config ignores non-array enabledFeatures""" + connection_manager = MagicMock() + connection_manager.conf = ServiceConfig( + endpoints=[], + last_updated_at=0, + blocked_uids=set(), + bypassed_ips=[], + received_any_stats=False, + ) + connection_manager.block = False + + res = { + "success": True, + "enabledFeatures": "not-an-array", + } + + update_service_config(connection_manager, res) + + assert connection_manager.conf.is_feature_enabled("not-an-array") is False + + def test_update_service_config_excluded_user_ids_not_array(): """Test that update_service_config ignores non-array excludedUserIdsFromRateLimiting""" connection_manager = MagicMock() diff --git a/aikido_zen/background_process/service_config.py b/aikido_zen/background_process/service_config.py index fc659cb7d..be34a9a81 100644 --- a/aikido_zen/background_process/service_config.py +++ b/aikido_zen/background_process/service_config.py @@ -25,6 +25,7 @@ def __init__( self.block_new_outgoing_requests = False self.outbound_domains = {} self.excluded_user_ids_from_rate_limiting = set() + self.enabled_features = set() def update( self, @@ -84,6 +85,14 @@ def is_user_excluded_from_rate_limiting(self, user_id): """Checks if the user ID is excluded from rate limiting""" return str(user_id) in self.excluded_user_ids_from_rate_limiting + def update_enabled_features(self, features): + """Sets the server-provided feature flags""" + self.enabled_features = set(features) + + def is_feature_enabled(self, feature) -> bool: + """Checks if `feature` is enabled via server-provided config""" + return feature in self.enabled_features + def update_outbound_domains(self, domains): self.outbound_domains = { domain["hostname"]: domain["mode"] for domain in domains diff --git a/aikido_zen/background_process/service_config_test.py b/aikido_zen/background_process/service_config_test.py index d4b7d38e0..dd57e90a6 100644 --- a/aikido_zen/background_process/service_config_test.py +++ b/aikido_zen/background_process/service_config_test.py @@ -321,6 +321,28 @@ def test_service_config_with_empty_allowlist(): assert len(admin_endpoint["allowedIPAddresses"]) == 0 +def test_service_config_enabled_features(): + """Test the update_enabled_features and is_feature_enabled methods""" + config = ServiceConfig( + endpoints=[], + last_updated_at=0, + blocked_uids=set(), + bypassed_ips=[], + received_any_stats=False, + ) + + # Initially empty + assert config.is_feature_enabled("realtime_updates") is False + + config.update_enabled_features(["realtime_updates"]) + assert config.is_feature_enabled("realtime_updates") is True + assert config.is_feature_enabled("some_other_feature") is False + + # Update replaces the set + config.update_enabled_features([]) + assert config.is_feature_enabled("realtime_updates") is False + + def test_excluded_user_ids_from_rate_limiting(): config = ServiceConfig( endpoints=[], From cb1d439acb06a38b088590b7a13d06b3419d28a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6ssler?= Date: Thu, 30 Jul 2026 15:29:17 +0200 Subject: [PATCH 2/2] Add more tests --- .../cloud_connection_manager/init_test.py | 55 +++++++++++++++++++ .../update_service_config_test.py | 46 ++++++++++++++++ 2 files changed, 101 insertions(+) diff --git a/aikido_zen/background_process/cloud_connection_manager/init_test.py b/aikido_zen/background_process/cloud_connection_manager/init_test.py index e14751904..f869e9bb0 100644 --- a/aikido_zen/background_process/cloud_connection_manager/init_test.py +++ b/aikido_zen/background_process/cloud_connection_manager/init_test.py @@ -99,4 +99,59 @@ def test_start_does_not_enable_sse_without_flag( mock_listen_for_config_updates.assert_not_called() +@patch("aikido_zen.background_process.cloud_connection_manager.on_start") +@patch( + "aikido_zen.background_process.cloud_connection_manager.start_polling_for_changes" +) +@patch( + "aikido_zen.background_process.cloud_connection_manager.send_heartbeats_every_x_secs" +) +@patch( + "aikido_zen.background_process.cloud_connection_manager.listen_for_config_updates" +) +def test_start_enables_sse_via_env_var( + mock_listen_for_config_updates, + mock_send_heartbeats, + mock_start_polling, + mock_on_start, + setup_cloud_connection_manager, + monkeypatch, +): + """SSE listening should still start via the AIKIDO_FEATURE_SSE env var, without the server flag""" + monkeypatch.setenv("AIKIDO_FEATURE_SSE", "true") + manager = setup_cloud_connection_manager + mock_on_start.return_value = {"success": True} + + manager.start(event_scheduler=MagicMock()) + + mock_listen_for_config_updates.assert_called_once() + + +@patch("aikido_zen.background_process.cloud_connection_manager.on_start") +@patch( + "aikido_zen.background_process.cloud_connection_manager.start_polling_for_changes" +) +@patch( + "aikido_zen.background_process.cloud_connection_manager.send_heartbeats_every_x_secs" +) +@patch( + "aikido_zen.background_process.cloud_connection_manager.listen_for_config_updates" +) +def test_start_does_not_enable_sse_for_unrelated_flag( + mock_listen_for_config_updates, + mock_send_heartbeats, + mock_start_polling, + mock_on_start, + setup_cloud_connection_manager, +): + """A server feature flag other than realtime_updates should not enable SSE""" + manager = setup_cloud_connection_manager + mock_on_start.return_value = {"success": True} + manager.conf.update_enabled_features(["some_other_feature"]) + + manager.start(event_scheduler=MagicMock()) + + mock_listen_for_config_updates.assert_not_called() + + # Additional tests can be added here for other edge cases or scenarios diff --git a/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py b/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py index e7c815bfa..104e44ec3 100644 --- a/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py +++ b/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py @@ -305,6 +305,52 @@ def test_update_service_config_enabled_features_not_array(): assert connection_manager.conf.is_feature_enabled("not-an-array") is False +def test_update_service_config_enabled_features_empty_list(): + """Test that update_service_config clears enabled features with an empty list""" + connection_manager = MagicMock() + connection_manager.conf = ServiceConfig( + endpoints=[], + last_updated_at=0, + blocked_uids=set(), + bypassed_ips=[], + received_any_stats=False, + ) + connection_manager.block = False + connection_manager.conf.update_enabled_features(["realtime_updates"]) + + res = { + "success": True, + "enabledFeatures": [], + } + + update_service_config(connection_manager, res) + + assert connection_manager.conf.is_feature_enabled("realtime_updates") is False + + +def test_update_service_config_enabled_features_not_updated_on_failure(): + """Test that enabledFeatures is left untouched when the response indicates failure""" + connection_manager = MagicMock() + connection_manager.conf = ServiceConfig( + endpoints=[], + last_updated_at=0, + blocked_uids=set(), + bypassed_ips=[], + received_any_stats=False, + ) + connection_manager.block = False + connection_manager.conf.update_enabled_features(["realtime_updates"]) + + res = { + "success": False, + "enabledFeatures": [], + } + + update_service_config(connection_manager, res) + + assert connection_manager.conf.is_feature_enabled("realtime_updates") is True + + def test_update_service_config_excluded_user_ids_not_array(): """Test that update_service_config ignores non-array excludedUserIdsFromRateLimiting""" connection_manager = MagicMock()