fix(core): make KeyAuth.close() wait for a running refresh

Author: mittalpk

core/src/stackit/core/auth_methods/key_auth.py

@@ -101,15 +101,18 @@ def close(self) -> None:
 
         Call this when the instance is no longer needed, for example after
         replacing it during key rotation in a long-running process. Without it
-        the refresh thread keeps running and keeps requesting tokens. Safe to
-        call more than once.
+        the refresh thread keeps running and keeps requesting tokens. When it
+        returns, no refresh is running and none will start. Safe to call more
+        than once.
         """
-        self._stop.set()
-        self.executor.shutdown(wait=False, cancel_futures=True)
+        # Set under the lock so the refresh loop cannot submit after shutdown.
+        with self.lock:
+            self._stop.set()
         atexit.unregister(self.__shutdown)
         thread = self._refresh_thread
         if thread is not None and thread is not threading.current_thread():
             thread.join()
+        self.executor.shutdown(wait=True, cancel_futures=True)
 
     def __create_initial_token(self) -> None:
         payload = {
@@ -147,6 +150,8 @@ def __start_token_refresh_task(self):
         def token_refresh_task():
             while not self._stop.wait(self.TOKEN_EXPIRY_CHECK_INTERVAL.total_seconds()):
                 with self.lock:
+                    if self._stop.is_set():
+                        return
                     if self.refresh_future is not None and self.refresh_future.done():
                         self.refresh_future.result()
                     if self.__is_token_expired(self.access_token) and (
@@ -175,6 +180,8 @@ def __refresh_token(self):
 
         last_exception = None
         for _ in range(self.MAX_REFRESH_RETRIES):
+            if self._stop.is_set():
+                return
             try:
                 response = requests.post(self.token_endpoint, data=body, timeout=self.timeout)
                 response.raise_for_status()

core/tests/core/test_auth.py

@@ -311,6 +311,7 @@ def test_call_propagates_token_refresh_failure(self, refresh_already_failed):
         auth.refresh_token = jwt.encode({"exp": 4_000_000_000}, "x" * 32, algorithm="HS256")
         auth.token_endpoint = KeyAuth.DEFAULT_TOKEN_ENDPOINT
         auth.refresh_future = None
+        auth._stop = Event()
         original_token = auth.access_token
         request = Request("GET", "https://example.com")
         refresh_error = requests.RequestException("refresh failed")
@@ -425,6 +426,35 @@ def set_initial_token(auth):
             time.sleep(0.1)  # ten refresh intervals
             assert mock_post.call_count == requests_after_close
 
+    def test_close_waits_for_a_refresh_in_progress(self, service_account_key_file_json):
+        def set_initial_token(auth):
+            auth.initial_token = "test-initial-token"
+
+        expired = jwt.encode({"exp": 1}, "secret", algorithm="HS256")
+        refresh_started = Event()
+        with (
+            patch.object(KeyAuth, "_KeyAuth__create_initial_token", new=set_initial_token),
+            patch.object(KeyAuth, "TOKEN_EXPIRY_CHECK_INTERVAL", timedelta(milliseconds=10)),
+            patch("requests.post") as mock_post,
+        ):
+            auth = self._make_auth(service_account_key_file_json, mock_post, expired)
+            initial_response = mock_post.return_value
+
+            def slow_refresh(*args, **kwargs):
+                refresh_started.set()
+                time.sleep(0.2)
+                return initial_response
+
+            mock_post.side_effect = slow_refresh
+            assert refresh_started.wait(timeout=5)
+
+            auth.close()
+
+            assert auth.refresh_future is not None and auth.refresh_future.done()
+            requests_after_close = mock_post.call_count
+            time.sleep(0.1)
+            assert mock_post.call_count == requests_after_close
+
     def test_close_is_idempotent(self, service_account_key_file_json):
         def set_initial_token(auth):
             auth.initial_token = "test-initial-token"