Fixed auth issues

azure-kusto-data/azure/kusto/data/env_utils.py

@@ -1,4 +1,6 @@
 import os
+from dataclasses import dataclass, astuple
+from typing import Optional
 
 
 def get_env(*args, optional=False, default=None):
@@ -19,19 +21,42 @@ def set_env(key, value):
 def get_app_id(optional=False):
     """Return the app id."""
     result = get_env("APP_ID", "AZURE_CLIENT_ID", optional=optional)
-    os.environ["AZURE_CLIENT_ID"] = result
+    if result:
+        set_env("AZURE_CLIENT_ID", result)
     return result
 
 
 def get_auth_id(optional=False):
     """Return the auth id."""
     result = get_env("AUTH_ID", "APP_AUTH_ID", "AZURE_TENANT_ID", optional=optional)
-    os.environ["AZURE_TENANT_ID"] = result
+    if result:
+        set_env("AZURE_TENANT_ID", result)
     return result
 
 
 def get_app_key(optional=False):
     """Return the app key."""
     result = get_env("APP_KEY", "AZURE_CLIENT_SECRET", optional=optional)
-    os.environ["AZURE_CLIENT_SECRET"] = result
+    if result:
+        set_env("AZURE_CLIENT_SECRET", result)
     return result
+
+
+@dataclass(frozen=True)
+class AppKeyAuth:
+    app_id: str
+    app_key: str
+    auth_id: str
+
+    def __iter__(self):
+        return iter(astuple(self))
+
+
+def prepare_app_key_auth(optional=False) -> Optional[AppKeyAuth]:
+    """Gets app key auth information from the env, sets the correct values for azidentity, and returns the AppKeyAuth object."""
+    app_id = get_app_id(optional=optional)
+    app_key = get_app_key(optional=optional)
+    auth_id = get_auth_id(optional=optional)
+    if app_id and app_key and auth_id:
+        return AppKeyAuth(app_id, app_key, auth_id)
+    return None

azure-kusto-data/tests/aio/test_async_token_providers.py

@@ -5,7 +5,7 @@
 
 from azure.kusto.data._decorators import aio_documented_by
 from azure.kusto.data._token_providers import *
-from azure.kusto.data.env_utils import get_env, get_app_id, get_auth_id, get_app_key
+from azure.kusto.data.env_utils import get_env, get_app_id, get_auth_id, get_app_key, prepare_app_key_auth
 from .test_kusto_client import run_aio_tests
 from ..test_token_providers import KUSTO_URI, TOKEN_VALUE, TEST_AZ_AUTH, TEST_MSI_AUTH, TEST_DEVICE_AUTH, TokenProviderTests, MockProvider
 
@@ -142,8 +142,7 @@ async def fail_callback():
     @pytest.mark.asyncio
     async def test_az_provider(self):
         if not TEST_AZ_AUTH:
-            print(" *** Skipped Az-Cli Provider Test ***")
-            return
+            pytest.skip(" *** Skipped Az-Cli Provider Test ***")
 
         print("Note!\nThe test 'test_az_provider' will fail if 'az login' was not called.")
         with AzCliTokenProvider(KUSTO_URI, is_async=True) as provider:
@@ -158,8 +157,7 @@ async def test_az_provider(self):
     @pytest.mark.asyncio
     async def test_msi_provider(self):
         if not TEST_MSI_AUTH:
-            print(" *** Skipped MSI Provider Test ***")
-            return
+            pytest.skip(" *** Skipped MSI Provider Test ***")
 
         user_msi_object_id = get_env("MSI_OBJECT_ID", optional=True)
         user_msi_client_id = get_env("MSI_CLIENT_ID", optional=True)
@@ -175,15 +173,15 @@ async def test_msi_provider(self):
                     token = await provider.get_token_async()
                     assert self.get_token_value(token) is not None
             else:
-                print(" *** Skipped MSI Provider Client Id Test ***")
+                pytest.skip(" *** Skipped MSI Provider Client Id Test ***")
 
             if user_msi_client_id is not None:
                 args = {"client_id": user_msi_client_id}
                 with MsiTokenProvider(KUSTO_URI, args, is_async=True) as provider:
                     token = await provider.get_token_async()
                     assert self.get_token_value(token) is not None
             else:
-                print(" *** Skipped MSI Provider Object Id Test ***")
+                pytest.skip(" *** Skipped MSI Provider Object Id Test ***")
 
     @aio_documented_by(TokenProviderTests.test_user_pass_provider)
     @pytest.mark.asyncio
@@ -201,14 +199,13 @@ async def test_user_pass_provider(self):
                 token = provider._get_token_from_cache_impl()
                 assert self.get_token_value(token) is not None
         else:
-            print(" *** Skipped User & Pass Provider Test ***")
+            pytest.skip(" *** Skipped User & Pass Provider Test ***")
 
     @aio_documented_by(TokenProviderTests.test_device_auth_provider)
     @pytest.mark.asyncio
     async def test_device_auth_provider(self):
         if not TEST_DEVICE_AUTH:
-            print(" *** Skipped User Device Flow Test ***")
-            return
+            pytest.skip(" *** Skipped User Device Flow Test ***")
 
         def callback(x, x2, x3):
             # break here if you debug this test, and get the code from 'x'
@@ -225,18 +222,14 @@ def callback(x, x2, x3):
     @aio_documented_by(TokenProviderTests.test_app_key_provider)
     @pytest.mark.asyncio
     async def test_app_key_provider(self):
-        # default details are for kusto-client-e2e-test-app
-        # to run the test, get the key from Azure portal
-        app_id = get_app_id(optional=True)
-        auth_id = get_auth_id(optional=True)
-        app_key = get_app_key(optional=True)
+        app_auth = prepare_app_key_auth(optional=True)
 
-        if app_id and app_key and auth_id:
-            with ApplicationKeyTokenProvider(KUSTO_URI, auth_id, app_id, app_key, is_async=True) as provider:
+        if app_auth:
+            with ApplicationKeyTokenProvider(KUSTO_URI, app_auth.auth_id, app_auth.app_id, app_auth.app_key, is_async=True) as provider:
                 token = await provider.get_token_async()
                 assert self.get_token_value(token) is not None
         else:
-            print(" *** Skipped App Id & Key Provider Test ***")
+            pytest.skip(" *** Skipped App Id & Key Provider Test ***")
 
     @aio_documented_by(TokenProviderTests.test_app_cert_provider)
     @pytest.mark.asyncio
@@ -269,10 +262,10 @@ async def test_app_cert_provider(self):
                         token = provider._get_token_from_cache_impl()
                         assert self.get_token_value(token) is not None
                 else:
-                    print(" *** Skipped App Cert SNI Provider Test ***")
+                    pytest.skip(" *** Skipped App Cert SNI Provider Test ***")
 
         else:
-            print(" *** Skipped App Cert Provider Test ***")
+            pytest.skip(" *** Skipped App Cert Provider Test ***")
 
     @aio_documented_by(TokenProviderTests.test_cloud_mfa_off)
     @pytest.mark.asyncio
@@ -342,9 +335,10 @@ async def inner():
     @aio_documented_by(TokenProviderTests.test_azure_identity_default_token_provider)
     @pytest.mark.asyncio
     async def test_azure_identity_token_provider(self):
-        app_id = get_app_id()
-        auth_id = get_auth_id()
-        app_key = get_app_key()
+        app_id, app_key, auth_id = prepare_app_key_auth(optional=True)
+        if not app_id or not app_key or not auth_id:
+            pytest.skip(" *** Skipped Azure Identity Provider Test ***")
+
         with AzureIdentityTokenCredentialProvider(KUSTO_URI, is_async=True, credential=AsyncDefaultAzureCredential()) as provider:
             token = await provider.get_token_async()
             assert TokenProviderTests.get_token_value(token) is not None

azure-kusto-data/tests/test_e2e_data.py

@@ -14,7 +14,7 @@
 import pytest
 from azure.identity import DefaultAzureCredential
 
-from azure.kusto.data.env_utils import get_env, get_app_id, get_auth_id, get_app_key, set_env
+from azure.kusto.data.env_utils import get_env, get_app_id, get_auth_id, get_app_key, set_env, prepare_app_key_auth
 from azure.kusto.data import KustoClient, KustoConnectionStringBuilder
 from azure.kusto.data._cloud_settings import CloudSettings, DEFAULT_DEV_KUSTO_SERVICE_RESOURCE_ID
 from azure.kusto.data._models import WellKnownDataSet
@@ -40,11 +40,10 @@ class TestE2E:
     test_streaming_data_csv_raw: ClassVar[str]
     engine_cs: ClassVar[Optional[str]]
     ai_engine_cs: ClassVar[Optional[str]]
-    app_id: ClassVar[Optional[str]]
-    app_key: ClassVar[Optional[str]]
-    auth_id: ClassVar[Optional[str]]
     test_db: ClassVar[Optional[str]]
     ai_test_db: ClassVar[Optional[str]]
+    cred: ClassVar[DefaultAzureCredential]
+    async_cred: ClassVar[AsyncDefaultAzureCredential]
 
     CHUNK_SIZE = 1024
 
@@ -104,25 +103,25 @@ def engine_kcsb_from_env(cls, app_insights=False, is_async=False) -> KustoConnec
         engine = cls.engine_cs if not app_insights else cls.ai_engine_cs
         return KustoConnectionStringBuilder.with_azure_token_credential(
             engine,
-            credential=DefaultAzureCredential(exclude_interactive_browser_credential=False)
-            if not is_async
-            else AsyncDefaultAzureCredential(exclude_interactive_browser_credential=False),
+            credential=cls.cred if not is_async else cls.async_cred,
         )
 
     @classmethod
     def setup_class(cls):
         cls.engine_cs = get_env("ENGINE_CONNECTION_STRING")
         cls.ai_engine_cs = get_env("APPLICATION_INSIGHTS_ENGINE_CONNECTION_STRING", optional=True)
 
-        cls.app_id = get_app_id()
-        cls.auth_id = get_auth_id()
-        cls.app_key = get_app_key()
+        # Called to set the env variables for the default azure credentials
+        prepare_app_key_auth(optional=True)
 
         set_env("AZURE_AUTHORITY_HOST", "login.microsoftonline.com")
 
         cls.test_db = get_env("TEST_DATABASE")
         cls.ai_test_db = get_env("APPLICATION_INSIGHTS_TEST_DATABASE", optional=True)  # name of e2e database could be changed
 
+        cls.cred = DefaultAzureCredential(exclude_interactive_browser_credential=False)
+        cls.async_cred = AsyncDefaultAzureCredential(exclude_interactive_browser_credential=False)
+
         cls.input_folder_path = cls.get_file_path()
 
         with open(os.path.join(cls.input_folder_path, "big_source.csv")) as f:
@@ -348,7 +347,7 @@ def test_cloud_info_404(self):
     @pytest.mark.parametrize("code", [301, 302, 307, 308])
     def test_no_redirects_fail_in_cloud(self, code):
         with KustoClient(
-            KustoConnectionStringBuilder.with_azure_token_credential(f"https://statusreturner.azurewebsites.net/{code}/nocloud", DefaultAzureCredential())
+            KustoConnectionStringBuilder.with_azure_token_credential(f"https://statusreturner.azurewebsites.net/{code}/nocloud", self.cred)
         ) as client:
             with pytest.raises(KustoServiceError) as ex:
                 client.execute("db", "table")
@@ -358,7 +357,7 @@ def test_no_redirects_fail_in_cloud(self, code):
     def test_no_redirects_fail_in_client(self, code):
         well_known_kusto_endpoints.add_trusted_hosts([MatchRule("statusreturner.azurewebsites.net", False)], False)
         with KustoClient(
-            KustoConnectionStringBuilder.with_azure_token_credential(f"https://statusreturner.azurewebsites.net/{code}/segment", DefaultAzureCredential())
+            KustoConnectionStringBuilder.with_azure_token_credential(f"https://statusreturner.azurewebsites.net/{code}/segment", self.cred)
         ) as client:
             with pytest.raises(KustoServiceError) as ex:
                 client.execute("db", "table")
@@ -368,7 +367,7 @@ def test_no_redirects_fail_in_client(self, code):
     @pytest.mark.parametrize("code", [301, 302, 307, 308])
     async def test_no_redirects_fail_in_cloud(self, code):
         async with AsyncKustoClient(
-            KustoConnectionStringBuilder.with_azure_token_credential(f"https://statusreturner.azurewebsites.net/{code}/nocloud", AsyncDefaultAzureCredential())
+            KustoConnectionStringBuilder.with_azure_token_credential(f"https://statusreturner.azurewebsites.net/{code}/nocloud", self.async_cred)
         ) as client:
             with pytest.raises(KustoServiceError) as ex:
                 await client.execute("db", "table")
@@ -379,7 +378,7 @@ async def test_no_redirects_fail_in_cloud(self, code):
     async def test_no_redirects_fail_in_client(self, code):
         well_known_kusto_endpoints.add_trusted_hosts([MatchRule("statusreturner.azurewebsites.net", False)], False)
         async with AsyncKustoClient(
-            KustoConnectionStringBuilder.with_azure_token_credential(f"https://statusreturner.azurewebsites.net/{code}/segment", AsyncDefaultAzureCredential())
+            KustoConnectionStringBuilder.with_azure_token_credential(f"https://statusreturner.azurewebsites.net/{code}/segment", self.async_cred)
         ) as client:
             with pytest.raises(KustoServiceError) as ex:
                 await client.execute("db", "table")

azure-kusto-data/tests/test_security.py

@@ -121,8 +121,7 @@ def test_user_app_token_auth():
 
 def test_interactive_login():
     if not TEST_INTERACTIVE_AUTH:
-        print(" *** Skipped interactive login Test ***")
-        return
+        pytest.skip(" *** Skipped interactive login Test ***")
 
     kcsb = KustoConnectionStringBuilder.with_interactive_login(KUSTO_TEST_URI)
     aad_helper = _AadHelper(kcsb, False)