feat contrib: aiogram 3

Relates: https://st.yandex-team.ru/, https://st.yandex-team.ru/

7 files changed, 344 insertions, 3 deletions

diff --git a/contrib/python/google-auth/py3/tests/oauth2/test_credentials.py b/contrib/python/google-auth/py3/tests/oauth2/test_credentials.py
index d6a1915862..5f1dcf3cbf 100644
--- a/contrib/python/google-auth/py3/tests/oauth2/test_credentials.py
+++ b/contrib/python/google-auth/py3/tests/oauth2/test_credentials.py
@@ -24,6 +24,7 @@ import pytest  # type: ignore
 from google.auth import _helpers
 from google.auth import exceptions
 from google.auth import transport
+from google.auth.credentials import TokenState
 from google.oauth2 import credentials
 
 
@@ -62,6 +63,7 @@ class TestCredentials(object):
         assert not credentials.expired
         # Scopes aren't required for these credentials
         assert not credentials.requires_scopes
+        assert credentials.token_state == TokenState.INVALID
         # Test properties
         assert credentials.refresh_token == self.REFRESH_TOKEN
         assert credentials.token_uri == self.TOKEN_URI
@@ -912,7 +914,11 @@ class TestCredentials(object):
         assert list(creds.__dict__).sort() == list(unpickled.__dict__).sort()
 
         for attr in list(creds.__dict__):
-            assert getattr(creds, attr) == getattr(unpickled, attr)
+            # Worker should always be None
+            if attr == "_refresh_worker":
+                assert getattr(unpickled, attr) is None
+            else:
+                assert getattr(creds, attr) == getattr(unpickled, attr)
 
     def test_pickle_and_unpickle_universe_domain(self):
         # old version of auth lib doesn't have _universe_domain, so the pickled
@@ -946,7 +952,7 @@ class TestCredentials(object):
         for attr in list(creds.__dict__):
             # For the _refresh_handler property, the unpickled creds should be
             # set to None.
-            if attr == "_refresh_handler":
+            if attr == "_refresh_handler" or attr == "_refresh_worker":
                 assert getattr(unpickled, attr) is None
             else:
                 assert getattr(creds, attr) == getattr(unpickled, attr)
@@ -958,6 +964,8 @@ class TestCredentials(object):
         # this mimics a pickle created with a previous class definition with
         # fewer attributes
         del creds.__dict__["_quota_project_id"]
+        del creds.__dict__["_refresh_handler"]
+        del creds.__dict__["_refresh_worker"]
 
         unpickled = pickle.loads(pickle.dumps(creds))
 
diff --git a/contrib/python/google-auth/py3/tests/test__refresh_worker.py b/contrib/python/google-auth/py3/tests/test__refresh_worker.py
new file mode 100644
index 0000000000..f842b02cac
--- /dev/null
+++ b/contrib/python/google-auth/py3/tests/test__refresh_worker.py
@@ -0,0 +1,156 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pickle
+import random
+import threading
+import time
+
+import mock
+import pytest  # type: ignore
+
+from google.auth import _refresh_worker, credentials, exceptions
+
+MAIN_THREAD_SLEEP_MS = 100 / 1000
+
+
+class MockCredentialsImpl(credentials.Credentials):
+    def __init__(self, sleep_seconds=None):
+        self.refresh_count = 0
+        self.token = None
+        self.sleep_seconds = sleep_seconds if sleep_seconds else None
+
+    def refresh(self, request):
+        if self.sleep_seconds:
+            time.sleep(self.sleep_seconds)
+        self.token = request
+        self.refresh_count += 1
+
+
+@pytest.fixture
+def test_thread_count():
+    return 25
+
+
+def _cred_spinlock(cred):
+    while cred.token is None:  # pragma: NO COVER
+        time.sleep(MAIN_THREAD_SLEEP_MS)
+
+
+def test_invalid_start_refresh():
+    w = _refresh_worker.RefreshThreadManager()
+    with pytest.raises(exceptions.InvalidValue):
+        w.start_refresh(None, None)
+
+
+def test_start_refresh():
+    w = _refresh_worker.RefreshThreadManager()
+    cred = MockCredentialsImpl()
+    request = mock.MagicMock()
+    assert w.start_refresh(cred, request)
+
+    assert w._worker is not None
+
+    _cred_spinlock(cred)
+
+    assert cred.token == request
+    assert cred.refresh_count == 1
+
+
+def test_nonblocking_start_refresh():
+    w = _refresh_worker.RefreshThreadManager()
+    cred = MockCredentialsImpl(sleep_seconds=1)
+    request = mock.MagicMock()
+    assert w.start_refresh(cred, request)
+
+    assert w._worker is not None
+    assert not cred.token
+    assert cred.refresh_count == 0
+
+
+def test_multiple_refreshes_multiple_workers(test_thread_count):
+    w = _refresh_worker.RefreshThreadManager()
+    cred = MockCredentialsImpl()
+    request = mock.MagicMock()
+
+    def _thread_refresh():
+        time.sleep(random.randrange(0, 5))
+        assert w.start_refresh(cred, request)
+
+    threads = [
+        threading.Thread(target=_thread_refresh) for _ in range(test_thread_count)
+    ]
+    for t in threads:
+        t.start()
+
+    _cred_spinlock(cred)
+
+    assert cred.token == request
+    # There is a chance only one thread has enough time to perform a refresh.
+    # Generally multiple threads will have time to perform a refresh
+    assert cred.refresh_count > 0
+
+
+def test_refresh_error():
+    w = _refresh_worker.RefreshThreadManager()
+    cred = mock.MagicMock()
+    request = mock.MagicMock()
+
+    cred.refresh.side_effect = exceptions.RefreshError("Failed to refresh")
+
+    assert w.start_refresh(cred, request)
+
+    while w._worker._error_info is None:  # pragma: NO COVER
+        time.sleep(MAIN_THREAD_SLEEP_MS)
+
+    assert w._worker is not None
+    assert isinstance(w._worker._error_info, exceptions.RefreshError)
+
+
+def test_refresh_error_call_refresh_again():
+    w = _refresh_worker.RefreshThreadManager()
+    cred = mock.MagicMock()
+    request = mock.MagicMock()
+
+    cred.refresh.side_effect = exceptions.RefreshError("Failed to refresh")
+
+    assert w.start_refresh(cred, request)
+
+    while w._worker._error_info is None:  # pragma: NO COVER
+        time.sleep(MAIN_THREAD_SLEEP_MS)
+
+    assert not w.start_refresh(cred, request)
+
+
+def test_refresh_dead_worker():
+    cred = MockCredentialsImpl()
+    request = mock.MagicMock()
+
+    w = _refresh_worker.RefreshThreadManager()
+    w._worker = None
+
+    w.start_refresh(cred, request)
+
+    _cred_spinlock(cred)
+
+    assert cred.token == request
+    assert cred.refresh_count == 1
+
+
+def test_pickle():
+    w = _refresh_worker.RefreshThreadManager()
+
+    pickled_manager = pickle.dumps(w)
+    manager = pickle.loads(pickled_manager)
+    assert isinstance(manager, _refresh_worker.RefreshThreadManager)
diff --git a/contrib/python/google-auth/py3/tests/test_credentials.py b/contrib/python/google-auth/py3/tests/test_credentials.py
index d64f3abb50..8e6bbc9633 100644
--- a/contrib/python/google-auth/py3/tests/test_credentials.py
+++ b/contrib/python/google-auth/py3/tests/test_credentials.py
@@ -14,6 +14,7 @@
 
 import datetime
 
+import mock
 import pytest  # type: ignore
 
 from google.auth import _helpers
@@ -23,6 +24,11 @@ from google.auth import credentials
 class CredentialsImpl(credentials.Credentials):
     def refresh(self, request):
         self.token = request
+        self.expiry = (
+            datetime.datetime.utcnow()
+            + _helpers.REFRESH_THRESHOLD
+            + datetime.timedelta(seconds=5)
+        )
 
     def with_quota_project(self, quota_project_id):
         raise NotImplementedError()
@@ -43,6 +49,13 @@ def test_credentials_constructor():
     assert not credentials.expired
     assert not credentials.valid
     assert credentials.universe_domain == "googleapis.com"
+    assert not credentials._use_non_blocking_refresh
+
+
+def test_with_non_blocking_refresh():
+    c = CredentialsImpl()
+    c.with_non_blocking_refresh()
+    assert c._use_non_blocking_refresh
 
 
 def test_expired_and_valid():
@@ -220,3 +233,108 @@ def test_create_scoped_if_required_not_scopes():
     )
 
     assert scoped_credentials is unscoped_credentials
+
+
+def test_nonblocking_refresh_fresh_credentials():
+    c = CredentialsImpl()
+
+    c._refresh_worker = mock.MagicMock()
+
+    request = "token"
+
+    c.refresh(request)
+    assert c.token_state == credentials.TokenState.FRESH
+
+    c.with_non_blocking_refresh()
+    c.before_request(request, "http://example.com", "GET", {})
+
+
+def test_nonblocking_refresh_invalid_credentials():
+    c = CredentialsImpl()
+    c.with_non_blocking_refresh()
+
+    request = "token"
+    headers = {}
+
+    assert c.token_state == credentials.TokenState.INVALID
+
+    c.before_request(request, "http://example.com", "GET", headers)
+    assert c.token_state == credentials.TokenState.FRESH
+    assert c.valid
+    assert c.token == "token"
+    assert headers["authorization"] == "Bearer token"
+    assert "x-identity-trust-boundary" not in headers
+
+
+def test_nonblocking_refresh_stale_credentials():
+    c = CredentialsImpl()
+    c.with_non_blocking_refresh()
+
+    request = "token"
+    headers = {}
+
+    # Invalid credentials MUST require a blocking refresh.
+    c.before_request(request, "http://example.com", "GET", headers)
+    assert c.token_state == credentials.TokenState.FRESH
+    assert not c._refresh_worker._worker
+
+    c.expiry = (
+        datetime.datetime.utcnow()
+        + _helpers.REFRESH_THRESHOLD
+        - datetime.timedelta(seconds=1)
+    )
+
+    # STALE credentials SHOULD spawn a non-blocking worker
+    assert c.token_state == credentials.TokenState.STALE
+    c.before_request(request, "http://example.com", "GET", headers)
+    assert c._refresh_worker._worker is not None
+
+    assert c.token_state == credentials.TokenState.FRESH
+    assert c.valid
+    assert c.token == "token"
+    assert headers["authorization"] == "Bearer token"
+    assert "x-identity-trust-boundary" not in headers
+
+
+def test_nonblocking_refresh_failed_credentials():
+    c = CredentialsImpl()
+    c.with_non_blocking_refresh()
+
+    request = "token"
+    headers = {}
+
+    # Invalid credentials MUST require a blocking refresh.
+    c.before_request(request, "http://example.com", "GET", headers)
+    assert c.token_state == credentials.TokenState.FRESH
+    assert not c._refresh_worker._worker
+
+    c.expiry = (
+        datetime.datetime.utcnow()
+        + _helpers.REFRESH_THRESHOLD
+        - datetime.timedelta(seconds=1)
+    )
+
+    # STALE credentials SHOULD spawn a non-blocking worker
+    assert c.token_state == credentials.TokenState.STALE
+    c._refresh_worker._worker = mock.MagicMock()
+    c._refresh_worker._worker._error_info = "Some Error"
+    c.before_request(request, "http://example.com", "GET", headers)
+    assert c._refresh_worker._worker is not None
+
+    assert c.token_state == credentials.TokenState.FRESH
+    assert c.valid
+    assert c.token == "token"
+    assert headers["authorization"] == "Bearer token"
+    assert "x-identity-trust-boundary" not in headers
+
+
+def test_token_state_no_expiry():
+    c = CredentialsImpl()
+
+    request = "token"
+    c.refresh(request)
+
+    c.expiry = None
+    assert c.token_state == credentials.TokenState.FRESH
+
+    c.before_request(request, "http://example.com", "GET", {})
diff --git a/contrib/python/google-auth/py3/tests/test_downscoped.py b/contrib/python/google-auth/py3/tests/test_downscoped.py
index b011380bdb..8cc2a30d16 100644
--- a/contrib/python/google-auth/py3/tests/test_downscoped.py
+++ b/contrib/python/google-auth/py3/tests/test_downscoped.py
@@ -25,6 +25,7 @@ from google.auth import credentials
 from google.auth import downscoped
 from google.auth import exceptions
 from google.auth import transport
+from google.auth.credentials import TokenState
 
 
 EXPRESSION = (
@@ -676,6 +677,7 @@ class TestCredentials(object):
 
         assert credentials.valid
         assert not credentials.expired
+        assert credentials.token_state == TokenState.FRESH
 
         credentials.before_request(request, "POST", "https://example.com/api", headers)
 
@@ -687,8 +689,24 @@ class TestCredentials(object):
 
         assert not credentials.valid
         assert credentials.expired
+        assert credentials.token_state == TokenState.STALE
 
         credentials.before_request(request, "POST", "https://example.com/api", headers)
+        assert credentials.token_state == TokenState.FRESH
+
+        # New token should be retrieved.
+        assert headers == {
+            "authorization": "Bearer {}".format(SUCCESS_RESPONSE["access_token"])
+        }
+
+        utcnow.return_value = datetime.datetime.min + datetime.timedelta(seconds=6000)
+
+        assert not credentials.valid
+        assert credentials.expired
+        assert credentials.token_state == TokenState.INVALID
+
+        credentials.before_request(request, "POST", "https://example.com/api", headers)
+        assert credentials.token_state == TokenState.FRESH
 
         # New token should be retrieved.
         assert headers == {
diff --git a/contrib/python/google-auth/py3/tests/test_external_account.py b/contrib/python/google-auth/py3/tests/test_external_account.py
index 5225dcf342..7f33b1dfa2 100644
--- a/contrib/python/google-auth/py3/tests/test_external_account.py
+++ b/contrib/python/google-auth/py3/tests/test_external_account.py
@@ -24,6 +24,7 @@ from google.auth import _helpers
 from google.auth import exceptions
 from google.auth import external_account
 from google.auth import transport
+from google.auth.credentials import TokenState
 
 
 IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE = (
@@ -1494,6 +1495,7 @@ class TestCredentials(object):
 
         assert credentials.valid
         assert not credentials.expired
+        assert credentials.token_state == TokenState.FRESH
 
         credentials.before_request(request, "POST", "https://example.com/api", headers)
 
@@ -1508,8 +1510,10 @@ class TestCredentials(object):
 
         assert not credentials.valid
         assert credentials.expired
+        assert credentials.token_state == TokenState.STALE
 
         credentials.before_request(request, "POST", "https://example.com/api", headers)
+        assert credentials.token_state == TokenState.FRESH
 
         # New token should be retrieved.
         assert headers == {
@@ -1551,8 +1555,10 @@ class TestCredentials(object):
 
         assert credentials.valid
         assert not credentials.expired
+        assert credentials.token_state == TokenState.FRESH
 
         credentials.before_request(request, "POST", "https://example.com/api", headers)
+        assert credentials.token_state == TokenState.FRESH
 
         # Cached token should be used.
         assert headers == {
@@ -1566,6 +1572,10 @@ class TestCredentials(object):
 
         assert not credentials.valid
         assert credentials.expired
+        assert credentials.token_state == TokenState.STALE
+
+        credentials.before_request(request, "POST", "https://example.com/api", headers)
+        assert credentials.token_state == TokenState.FRESH
 
         credentials.before_request(request, "POST", "https://example.com/api", headers)
 
diff --git a/contrib/python/google-auth/py3/tests/test_external_account_authorized_user.py b/contrib/python/google-auth/py3/tests/test_external_account_authorized_user.py
index 7ffd5078c8..7213a23486 100644
--- a/contrib/python/google-auth/py3/tests/test_external_account_authorized_user.py
+++ b/contrib/python/google-auth/py3/tests/test_external_account_authorized_user.py
@@ -44,6 +44,8 @@ CLIENT_SECRET = "password"
 BASIC_AUTH_ENCODING = "dXNlcm5hbWU6cGFzc3dvcmQ="
 SCOPES = ["email", "profile"]
 NOW = datetime.datetime(1990, 8, 27, 6, 54, 30)
+FAKE_UNIVERSE_DOMAIN = "fake-universe-domain"
+DEFAULT_UNIVERSE_DOMAIN = external_account_authorized_user._DEFAULT_UNIVERSE_DOMAIN
 
 
 class TestCredentials(object):
@@ -98,6 +100,7 @@ class TestCredentials(object):
         assert creds.refresh_token == REFRESH_TOKEN
         assert creds.audience == AUDIENCE
         assert creds.token_url == TOKEN_URL
+        assert creds.universe_domain == DEFAULT_UNIVERSE_DOMAIN
 
     def test_basic_create(self):
         creds = external_account_authorized_user.Credentials(
@@ -105,6 +108,7 @@ class TestCredentials(object):
             expiry=datetime.datetime.max,
             scopes=SCOPES,
             revoke_url=REVOKE_URL,
+            universe_domain=FAKE_UNIVERSE_DOMAIN,
         )
 
         assert creds.expiry == datetime.datetime.max
@@ -115,6 +119,7 @@ class TestCredentials(object):
         assert creds.scopes == SCOPES
         assert creds.is_user
         assert creds.revoke_url == REVOKE_URL
+        assert creds.universe_domain == FAKE_UNIVERSE_DOMAIN
 
     def test_stunted_create_no_refresh_token(self):
         with pytest.raises(ValueError) as excinfo:
@@ -339,6 +344,7 @@ class TestCredentials(object):
         assert info["token_info_url"] == TOKEN_INFO_URL
         assert info["client_id"] == CLIENT_ID
         assert info["client_secret"] == CLIENT_SECRET
+        assert info["universe_domain"] == DEFAULT_UNIVERSE_DOMAIN
         assert "token" not in info
         assert "expiry" not in info
         assert "revoke_url" not in info
@@ -350,6 +356,7 @@ class TestCredentials(object):
             expiry=NOW,
             revoke_url=REVOKE_URL,
             quota_project_id=QUOTA_PROJECT_ID,
+            universe_domain=FAKE_UNIVERSE_DOMAIN,
         )
         info = creds.info
 
@@ -363,6 +370,7 @@ class TestCredentials(object):
         assert info["expiry"] == NOW.isoformat() + "Z"
         assert info["revoke_url"] == REVOKE_URL
         assert info["quota_project_id"] == QUOTA_PROJECT_ID
+        assert info["universe_domain"] == FAKE_UNIVERSE_DOMAIN
 
     def test_to_json(self):
         creds = self.make_credentials()
@@ -375,6 +383,7 @@ class TestCredentials(object):
         assert info["token_info_url"] == TOKEN_INFO_URL
         assert info["client_id"] == CLIENT_ID
         assert info["client_secret"] == CLIENT_SECRET
+        assert info["universe_domain"] == DEFAULT_UNIVERSE_DOMAIN
         assert "token" not in info
         assert "expiry" not in info
         assert "revoke_url" not in info
@@ -386,6 +395,7 @@ class TestCredentials(object):
             expiry=NOW,
             revoke_url=REVOKE_URL,
             quota_project_id=QUOTA_PROJECT_ID,
+            universe_domain=FAKE_UNIVERSE_DOMAIN,
         )
         json_info = creds.to_json()
         info = json.loads(json_info)
@@ -400,6 +410,7 @@ class TestCredentials(object):
         assert info["expiry"] == NOW.isoformat() + "Z"
         assert info["revoke_url"] == REVOKE_URL
         assert info["quota_project_id"] == QUOTA_PROJECT_ID
+        assert info["universe_domain"] == FAKE_UNIVERSE_DOMAIN
 
     def test_to_json_full_with_strip(self):
         creds = self.make_credentials(
@@ -467,6 +478,26 @@ class TestCredentials(object):
         assert new_creds._revoke_url == creds._revoke_url
         assert new_creds._quota_project_id == creds._quota_project_id
 
+    def test_with_universe_domain(self):
+        creds = self.make_credentials(
+            token=ACCESS_TOKEN,
+            expiry=NOW,
+            revoke_url=REVOKE_URL,
+            quota_project_id=QUOTA_PROJECT_ID,
+        )
+        new_creds = creds.with_universe_domain(FAKE_UNIVERSE_DOMAIN)
+        assert new_creds._audience == creds._audience
+        assert new_creds._refresh_token == creds._refresh_token
+        assert new_creds._token_url == creds._token_url
+        assert new_creds._token_info_url == creds._token_info_url
+        assert new_creds._client_id == creds._client_id
+        assert new_creds._client_secret == creds._client_secret
+        assert new_creds.token == creds.token
+        assert new_creds.expiry == creds.expiry
+        assert new_creds._revoke_url == creds._revoke_url
+        assert new_creds._quota_project_id == QUOTA_PROJECT_ID
+        assert new_creds.universe_domain == FAKE_UNIVERSE_DOMAIN
+
     def test_from_file_required_options_only(self, tmpdir):
         from_creds = self.make_credentials()
         config_file = tmpdir.join("config.json")
diff --git a/contrib/python/google-auth/py3/tests/test_impersonated_credentials.py b/contrib/python/google-auth/py3/tests/test_impersonated_credentials.py
index 9696e823ff..7295bba429 100644
--- a/contrib/python/google-auth/py3/tests/test_impersonated_credentials.py
+++ b/contrib/python/google-auth/py3/tests/test_impersonated_credentials.py
@@ -243,7 +243,7 @@ class TestImpersonatedCredentials(object):
         request_kwargs = request.call_args[1]
         assert request_kwargs["url"] == self.IAM_ENDPOINT_OVERRIDE
 
-    @pytest.mark.parametrize("time_skew", [100, -100])
+    @pytest.mark.parametrize("time_skew", [150, -150])
     def test_refresh_source_credentials(self, time_skew):
         credentials = self.make_credentials(lifetime=None)