From 06e693e0f995708f7140a48308c91bce7b0476c6 Mon Sep 17 00:00:00 2001 From: tyiuhc Date: Tue, 6 Aug 2024 13:58:53 -0700 Subject: [PATCH] update cohort_sync_config fields: include polling and remove request delay, use enum for serverzone, update tests accordingly --- .../cohort/cohort_download_api.py | 13 +++++++------ .../cohort/cohort_sync_config.py | 7 ++++--- .../deployment/deployment_runner.py | 4 +--- src/amplitude_experiment/local/client.py | 1 - tests/cohort/cohort_download_api_test.py | 2 +- tests/deployment/deployment_runner_test.py | 5 +++-- tests/local/client_eu_test.py | 3 +-- tests/local/client_test.py | 3 +-- 8 files changed, 18 insertions(+), 20 deletions(-) diff --git a/src/amplitude_experiment/cohort/cohort_download_api.py b/src/amplitude_experiment/cohort/cohort_download_api.py index 6f9aae0..65da244 100644 --- a/src/amplitude_experiment/cohort/cohort_download_api.py +++ b/src/amplitude_experiment/cohort/cohort_download_api.py @@ -9,6 +9,8 @@ from ..connection_pool import HTTPConnectionPool from ..exception import HTTPErrorResponseException, CohortTooLargeException +COHORT_REQUEST_RETRY_DELAY_MILLIS = 100 + class CohortDownloadApi: @@ -17,13 +19,11 @@ def get_cohort(self, cohort_id: str, cohort: Optional[Cohort]) -> Cohort: class DirectCohortDownloadApi(CohortDownloadApi): - def __init__(self, api_key: str, secret_key: str, max_cohort_size: int, cohort_request_delay_millis: int, - server_url: str, logger: logging.Logger): + def __init__(self, api_key: str, secret_key: str, max_cohort_size: int, server_url: str, logger: logging.Logger): super().__init__() self.api_key = api_key self.secret_key = secret_key self.max_cohort_size = max_cohort_size - self.cohort_request_delay_millis = cohort_request_delay_millis self.server_url = server_url self.logger = logger self.__setup_connection_pool() @@ -48,10 +48,11 @@ def get_cohort(self, cohort_id: str, cohort: Optional[Cohort]) -> Cohort or None group_type=cohort_info['groupType'], ) elif response.status == 204: - self.logger.debug(f"getCohortMembers({cohort_id}): Cohort not modified" ) + self.logger.debug(f"getCohortMembers({cohort_id}): Cohort not modified") return elif response.status == 413: - raise CohortTooLargeException(f"Cohort exceeds max cohort size of {self.max_cohort_size}: {response.status}") + raise CohortTooLargeException( + f"Cohort exceeds max cohort size of {self.max_cohort_size}: {response.status}") elif response.status != 202: raise HTTPErrorResponseException(response.status, f"Unexpected response code: {response.status}") @@ -61,7 +62,7 @@ def get_cohort(self, cohort_id: str, cohort: Optional[Cohort]) -> Cohort or None self.logger.debug(f"getCohortMembers({cohort_id}): request-status error {errors} - {e}") if errors >= 3 or isinstance(e, CohortTooLargeException): raise e - time.sleep(self.cohort_request_delay_millis/1000) + time.sleep(COHORT_REQUEST_RETRY_DELAY_MILLIS / 1000) def _get_cohort_members_request(self, cohort_id: str, last_modified: int) -> HTTPResponse: headers = { diff --git a/src/amplitude_experiment/cohort/cohort_sync_config.py b/src/amplitude_experiment/cohort/cohort_sync_config.py index 0ee25b2..e3b0090 100644 --- a/src/amplitude_experiment/cohort/cohort_sync_config.py +++ b/src/amplitude_experiment/cohort/cohort_sync_config.py @@ -10,14 +10,15 @@ class CohortSyncConfig: api_key (str): The project API Key secret_key (str): The project Secret Key max_cohort_size (int): The maximum cohort size that can be downloaded - cohort_request_delay_millis (int): The delay in milliseconds between cohort download requests + cohort_polling_interval_millis (int): The interval, in milliseconds, at which to poll for + cohort updates, minimum 60000 cohort_server_url (str): The server endpoint from which to request cohorts """ def __init__(self, api_key: str, secret_key: str, max_cohort_size: int = 2147483647, - cohort_request_delay_millis: int = 5000, cohort_server_url: str = DEFAULT_COHORT_SYNC_URL): + cohort_polling_interval_millis: int = 60000, cohort_server_url: str = DEFAULT_COHORT_SYNC_URL): self.api_key = api_key self.secret_key = secret_key self.max_cohort_size = max_cohort_size - self.cohort_request_delay_millis = cohort_request_delay_millis + self.cohort_polling_interval_millis = max(cohort_polling_interval_millis, 60000) self.cohort_server_url = cohort_server_url diff --git a/src/amplitude_experiment/deployment/deployment_runner.py b/src/amplitude_experiment/deployment/deployment_runner.py index f61b657..aa8aa64 100644 --- a/src/amplitude_experiment/deployment/deployment_runner.py +++ b/src/amplitude_experiment/deployment/deployment_runner.py @@ -10,8 +10,6 @@ from ..local.poller import Poller from ..util.flag_config import get_all_cohort_ids_from_flag, get_all_cohort_ids_from_flags -COHORT_POLLING_INTERVAL_MILLIS = 60000 - class DeploymentRunner: def __init__( @@ -31,7 +29,7 @@ def __init__( self.lock = threading.Lock() self.flag_poller = Poller(self.config.flag_config_polling_interval_millis / 1000, self.__periodic_flag_update) if self.cohort_loader: - self.cohort_poller = Poller(COHORT_POLLING_INTERVAL_MILLIS / 1000, + self.cohort_poller = Poller(self.config.cohort_sync_config.cohort_polling_interval_millis / 1000, self.__update_cohorts) self.logger = logger diff --git a/src/amplitude_experiment/local/client.py b/src/amplitude_experiment/local/client.py index 147de5d..ba917d4 100644 --- a/src/amplitude_experiment/local/client.py +++ b/src/amplitude_experiment/local/client.py @@ -61,7 +61,6 @@ def __init__(self, api_key: str, config: LocalEvaluationConfig = None): cohort_download_api = DirectCohortDownloadApi(self.config.cohort_sync_config.api_key, self.config.cohort_sync_config.secret_key, self.config.cohort_sync_config.max_cohort_size, - self.config.cohort_sync_config.cohort_request_delay_millis, self.config.cohort_sync_config.cohort_server_url, self.logger) diff --git a/tests/cohort/cohort_download_api_test.py b/tests/cohort/cohort_download_api_test.py index e8b856e..2a42394 100644 --- a/tests/cohort/cohort_download_api_test.py +++ b/tests/cohort/cohort_download_api_test.py @@ -19,7 +19,7 @@ def response(code: int, body: dict = None): class CohortDownloadApiTest(unittest.TestCase): def setUp(self): - self.api = DirectCohortDownloadApi('api', 'secret', 15000, 100, "https://example.amplitude.com", mock.create_autospec(logging.Logger)) + self.api = DirectCohortDownloadApi('api', 'secret', 15000, "https://example.amplitude.com", mock.create_autospec(logging.Logger)) def test_cohort_download_success(self): cohort = Cohort(id="1234", last_modified=0, size=1, member_ids={'user'}) diff --git a/tests/deployment/deployment_runner_test.py b/tests/deployment/deployment_runner_test.py index 447c190..f64a332 100644 --- a/tests/deployment/deployment_runner_test.py +++ b/tests/deployment/deployment_runner_test.py @@ -5,6 +5,7 @@ from src.amplitude_experiment import LocalEvaluationConfig from src.amplitude_experiment.cohort.cohort_loader import CohortLoader +from src.amplitude_experiment.cohort.cohort_sync_config import CohortSyncConfig from src.amplitude_experiment.flag.flag_config_api import FlagConfigApi from src.amplitude_experiment.deployment.deployment_runner import DeploymentRunner @@ -41,7 +42,7 @@ def test_start_throws_if_first_flag_config_load_fails(self): logger = mock.create_autospec(logging.Logger) cohort_loader = CohortLoader(cohort_download_api, cohort_storage) runner = DeploymentRunner( - LocalEvaluationConfig(), + LocalEvaluationConfig(cohort_sync_config=CohortSyncConfig('api_key', 'secret_key')), flag_api, flag_config_storage, cohort_storage, @@ -61,7 +62,7 @@ def test_start_does_not_throw_if_cohort_load_fails(self): logger = mock.create_autospec(logging.Logger) cohort_loader = CohortLoader(cohort_download_api, cohort_storage) runner = DeploymentRunner( - LocalEvaluationConfig(), + LocalEvaluationConfig(cohort_sync_config=CohortSyncConfig('api_key', 'secret_key')), flag_api, flag_config_storage, cohort_storage, logger, diff --git a/tests/local/client_eu_test.py b/tests/local/client_eu_test.py index ed8b3af..91e792d 100644 --- a/tests/local/client_eu_test.py +++ b/tests/local/client_eu_test.py @@ -17,8 +17,7 @@ def setUpClass(cls) -> None: api_key = os.environ['EU_API_KEY'] secret_key = os.environ['EU_SECRET_KEY'] cohort_sync_config = CohortSyncConfig(api_key=api_key, - secret_key=secret_key, - cohort_request_delay_millis=100) + secret_key=secret_key) cls._local_evaluation_client = ( LocalEvaluationClient(SERVER_API_KEY, LocalEvaluationConfig(debug=False, server_zone=ServerZone.EU, cohort_sync_config=cohort_sync_config))) diff --git a/tests/local/client_test.py b/tests/local/client_test.py index 87085c6..b6c50eb 100644 --- a/tests/local/client_test.py +++ b/tests/local/client_test.py @@ -22,8 +22,7 @@ def setUpClass(cls) -> None: api_key = os.environ['API_KEY'] secret_key = os.environ['SECRET_KEY'] cohort_sync_config = CohortSyncConfig(api_key=api_key, - secret_key=secret_key, - cohort_request_delay_millis=100) + secret_key=secret_key) cls._local_evaluation_client = ( LocalEvaluationClient(SERVER_API_KEY, LocalEvaluationConfig(debug=False, cohort_sync_config=cohort_sync_config)))