add oauth string auth helper

yandex-cloud · Dec 10, 2024 · 4824e62 · 4824e62
1 parent 5b73634
commit 4824e62
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/src/yandex_cloud_ml_sdk/_auth.py b/src/yandex_cloud_ml_sdk/_auth.py
@@ -335,6 +335,7 @@ async def get_auth_provider(
 ) -> BaseAuth:
     simple_iam_regexp = re.compile(r'^t\d\.')
     iam_regexp = re.compile(r't1\.[A-Z0-9a-z_-]+[=]{0,2}\.[A-Z0-9a-z_-]{86}[=]{0,2}')
+    simple_oauth_regexp = re.compile(r'y[0123]_[-\w]')
 
     result: BaseAuth | None = None
     if isinstance(auth, str):
@@ -347,6 +348,8 @@ async def get_auth_provider(
                     UserWarning,
                     stacklevel=2,
                 )
+        elif simple_oauth_regexp.match(auth):
+            result = OAuthTokenAuth(auth)
         else:
             result = APIKeyAuth(auth)
     elif isinstance(auth, BaseAuth):

diff --git a/tests/auth/test_get_auth_provider.py b/tests/auth/test_get_auth_provider.py
@@ -46,6 +46,14 @@ async def test_input_data(monkeypatch):
     auth = await get_auth_provider(auth=NoAuth(), endpoint=None, yc_profile=None)
     assert isinstance(auth, NoAuth)
 
+    with pytest.warns(UserWarning, match=r"Sharing your personal OAuth token is not safe"):
+        auth = await get_auth_provider(
+            auth="y3_ABC-_-abc123",
+            endpoint=None,
+            yc_profile=None
+        )
+    assert isinstance(auth, OAuthTokenAuth)
+
 
 @pytest.mark.filterwarnings(r"ignore:.*OAuth:UserWarning")
 async def test_order(monkeypatch, mock_client, process_maker):