fix: flush pending assignment and exposure events on LocalEvaluationClient.stop()

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

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()