diff --git a/src/amplitude_experiment/local/client.py b/src/amplitude_experiment/local/client.py index 4e87eda..b671549 100644 --- a/src/amplitude_experiment/local/client.py +++ b/src/amplitude_experiment/local/client.py @@ -1,5 +1,6 @@ +from concurrent.futures import wait from threading import Lock -from typing import Any, List, Dict, Set +from typing import Any, List, Dict, Set, Optional from amplitude import Amplitude @@ -158,12 +159,35 @@ def __setup_connection_pool(self): self._connection_pool = HTTPConnectionPool(host, max_size=1, idle_timeout=30, read_timeout=timeout, scheme=scheme) - def stop(self) -> None: + def stop(self, timeout: Optional[float] = 10.0) -> None: """ - Stop polling for flag configurations. Close resource like connection pool with client + Stop polling for flag configurations, flush pending assignment and exposure events, and close resources + like the connection pool. + + Parameters: + timeout (float | None): Maximum time, in seconds, to wait for pending assignment and exposure + events to finish sending before returning. Defaults to 10 seconds. Pass None to wait + indefinitely. """ self.deployment_runner.stop() self._connection_pool.close() + self.__shutdown_event_services(timeout) + + def __shutdown_event_services(self, timeout: Optional[float]) -> None: + instances = [service.amplitude for service in (self.assignment_service, self.exposure_service) + if service is not None] + if not instances: + return + futures = [] + for instance in instances: + futures.extend(f for f in (instance.flush() or []) if f is not None) + if futures: + _, not_done = wait(futures, timeout=timeout) + if not_done: + self.logger.warning(f"[Experiment] Stop timed out after {timeout}s waiting for " + f"{len(not_done)} pending event batch(es) to flush") + for instance in instances: + instance.shutdown() def __enter__(self) -> 'LocalEvaluationClient': return self diff --git a/tests/local/stop_flush_test.py b/tests/local/stop_flush_test.py new file mode 100644 index 0000000..0e791a2 --- /dev/null +++ b/tests/local/stop_flush_test.py @@ -0,0 +1,78 @@ +import time +import unittest +from concurrent.futures import Future +from unittest.mock import MagicMock + +from src.amplitude_experiment import LocalEvaluationClient, LocalEvaluationConfig +from src.amplitude_experiment.assignment import AssignmentConfig +from src.amplitude_experiment.exposure.exposure_config import ExposureConfig + +API_KEY = 'server-api-key' + + +def completed_future() -> Future: + future = Future() + future.set_result(None) + return future + + +class LocalEvaluationClientStopTestCase(unittest.TestCase): + + def _client_with_event_services(self) -> LocalEvaluationClient: + config = LocalEvaluationConfig( + assignment_config=AssignmentConfig(api_key='analytics-api-key'), + exposure_config=ExposureConfig(api_key='analytics-api-key'), + ) + return LocalEvaluationClient(API_KEY, config) + + def test_stop_flushes_then_shuts_down_assignment_and_exposure(self): + client = self._client_with_event_services() + assignment_amplitude = MagicMock() + assignment_amplitude.flush.return_value = [completed_future()] + exposure_amplitude = MagicMock() + exposure_amplitude.flush.return_value = [None] + client.assignment_service.amplitude = assignment_amplitude + client.exposure_service.amplitude = exposure_amplitude + + client.stop() + + assignment_amplitude.flush.assert_called_once() + exposure_amplitude.flush.assert_called_once() + assignment_amplitude.shutdown.assert_called_once() + exposure_amplitude.shutdown.assert_called_once() + + def test_stop_timeout_bounds_wait_on_pending_events(self): + client = self._client_with_event_services() + never_completes = Future() + assignment_amplitude = MagicMock() + assignment_amplitude.flush.return_value = [never_completes] + client.assignment_service.amplitude = assignment_amplitude + client.exposure_service.amplitude = MagicMock(flush=MagicMock(return_value=[])) + + start = time.monotonic() + client.stop(timeout=0.2) + elapsed = time.monotonic() - start + + self.assertLess(elapsed, 2) + assignment_amplitude.shutdown.assert_called_once() + + def test_stop_without_event_services(self): + client = LocalEvaluationClient(API_KEY, LocalEvaluationConfig()) + client.stop() + + def test_context_manager_exit_flushes(self): + client = self._client_with_event_services() + exposure_amplitude = MagicMock() + exposure_amplitude.flush.return_value = [completed_future()] + client.assignment_service.amplitude = MagicMock(flush=MagicMock(return_value=[])) + client.exposure_service.amplitude = exposure_amplitude + + with client: + pass + + exposure_amplitude.flush.assert_called_once() + exposure_amplitude.shutdown.assert_called_once() + + +if __name__ == '__main__': + unittest.main()