fix SSO token refresh bug (#1030)

.gitignore

@@ -17,6 +17,7 @@ eggs/
 lib/
 lib64/
 parts/
+.run/
 sdist/
 var/
 *.egg-info/

CHANGES.rst

@@ -1,9 +1,11 @@
 Changes
 -------
-2.5.3 (TBD)
+2.5.3 (2023-08-06)
 ^^^^^^^^^^^^^^^^^^
 * add more support for Python 3.11
 * bump botocore to 1.31.17
+* add waiter.wait return
+* fix SSO token refresh bug #1025
 
 2.5.2 (2023-07-06)
 ^^^^^^^^^^^^^^^^^^

aiobotocore/tokens.py

@@ -83,12 +83,13 @@ async def _protected_refresh(self):
 
 class AioSSOTokenProvider(SSOTokenProvider):
     async def _attempt_create_token(self, token):
-        response = await self._client.create_token(
-            grantType=self._GRANT_TYPE,
-            clientId=token["clientId"],
-            clientSecret=token["clientSecret"],
-            refreshToken=token["refreshToken"],
-        )
+        async with self._client as client:
+            response = await client.create_token(
+                grantType=self._GRANT_TYPE,
+                clientId=token["clientId"],
+                clientSecret=token["clientSecret"],
+                refreshToken=token["refreshToken"],
+            )
         expires_in = timedelta(seconds=response["expiresIn"])
         new_token = {
             "startUrl": self._sso_config["sso_start_url"],

aiobotocore/utils.py

@@ -480,7 +480,8 @@ async def get_bucket_region(self, bucket, response):
 
         # Finally, HEAD the bucket. No other choice sadly.
         try:
-            response = await self._client.head_bucket(Bucket=bucket)
+            async with self._client as client:
+                response = await client.head_bucket(Bucket=bucket)
             headers = response['ResponseMetadata']['HTTPHeaders']
         except ClientError as e:
             headers = e.response['ResponseMetadata']['HTTPHeaders']
@@ -595,7 +596,8 @@ async def get_bucket_region(self, bucket, response):
 
         # Finally, HEAD the bucket. No other choice sadly.
         try:
-            response = await self._client.head_bucket(Bucket=bucket)
+            async with self._client as client:
+                response = await client.head_bucket(Bucket=bucket)
             headers = response['ResponseMetadata']['HTTPHeaders']
         except ClientError as e:
             headers = e.response['ResponseMetadata']['HTTPHeaders']

aiobotocore/waiter.py

@@ -46,7 +46,7 @@ def create_waiter_with_client(waiter_name, waiter_model, client):
     # Waiter.wait method. This is needed to attach a docstring to the
     # method.
     async def wait(self, **kwargs):
-        await AIOWaiter.wait(self, **kwargs)
+        return await AIOWaiter.wait(self, **kwargs)
 
     wait.__doc__ = WaiterDocstring(
         waiter_name=waiter_name,
@@ -118,7 +118,7 @@ async def wait(self, **kwargs):
                 logger.debug(
                     "Waiting complete, waiter matched the " "success state."
                 )
-                return
+                return response
             if current_state == 'failure':
                 reason = 'Waiter encountered a terminal failure state: %s' % (
                     acceptor.explanation

pytest.ini

@@ -1,4 +1,5 @@
 [pytest]
+cache_dir = /tmp/pytest_aiobotocore_cache
 markers =
     moto
     config_kwargs

tests/python3.8/boto_tests/test_tokens.py

@@ -296,7 +296,9 @@ async def test_sso_token_provider_refresh(test_case):
         token_cache[cache_key] = cached_token
 
     mock_session = _create_mock_session(config)
-    mock_sso_oidc = mock.Mock()
+    mock_sso_oidc = mock.AsyncMock()
+    mock_sso_oidc.__aenter__.return_value = mock_sso_oidc
+    mock_sso_oidc.__aexit__.return_value = None
     mock_session.create_client.return_value = mock_sso_oidc
 
     refresh_response = test_case.pop("refreshResponse", None)

tests/test_config.py

@@ -12,7 +12,8 @@
 
 # NOTE: this doesn't require moto but needs to be marked to run with coverage
 @pytest.mark.moto
-def test_connector_args():
+@pytest.mark.asyncio
+async def test_connector_args():
     with pytest.raises(ParamValidationError):
         # wrong type
         connector_args = dict(use_dns_cache=1)