DRY in env code, fix e2e tests

.github/workflows/build.yml

@@ -41,6 +41,7 @@ jobs:
           APP_ID: ${{ secrets.APP_ID }}
           APP_KEY: ${{ secrets.APP_KEY }}
           AUTH_ID: ${{ secrets.AUTH_ID }}
+          ENGINE_CONNECTION_STRING: ${{ secrets.ENGINE_CONNECTION_STRING }}
       - name: EtoE Test with pytest
         env:
           APP_ID: ${{ secrets.APP_ID }}

CHANGELOG.md

@@ -3,6 +3,10 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+## Unreleased
+### Fixed
+- Internal fixes for environment variables
+
 ## [4.2.0] - 2023-04-16
 ### Added
 - Added Initial Catalog (Default Database) parameter to ConnectionStringBuilder

azure-kusto-data/azure/kusto/data/_cloud_settings.py

@@ -1,5 +1,4 @@
 import dataclasses
-import os
 from threading import Lock
 from typing import Optional, Dict
 from urllib.parse import urljoin
@@ -9,6 +8,7 @@
 from azure.core.tracing.decorator import distributed_trace
 from azure.core.tracing import SpanKind
 
+from .env_utils import get_env
 from ._telemetry import Span, MonitoredActivity
 from .exceptions import KustoServiceError
 
@@ -46,7 +46,7 @@ class CloudSettings:
     _cloud_cache_lock = Lock()
 
     DEFAULT_CLOUD = CloudInfo(
-        login_endpoint=os.environ.get(DEFAULT_AUTH_ENV_VAR_NAME, DEFAULT_PUBLIC_LOGIN_URL),
+        login_endpoint=get_env(DEFAULT_AUTH_ENV_VAR_NAME, default=DEFAULT_PUBLIC_LOGIN_URL),
         login_mfa_required=False,
         kusto_client_app_id=DEFAULT_KUSTO_CLIENT_APP_ID,
         kusto_client_redirect_uri=DEFAULT_REDIRECT_URI,

azure-kusto-data/azure/kusto/data/_token_providers.py

@@ -564,9 +564,8 @@ def _get_token_impl(self) -> Optional[dict]:
         token = self._msal_client.acquire_token_for_client(scopes=self._scopes)
         return self._valid_token_or_throw(token)
 
-    def _get_token_from_cache_impl(self) -> dict:
-        token = self._msal_client.acquire_token_silent(scopes=self._scopes, account=None)
-        return self._valid_token_or_none(token)
+    def _get_token_from_cache_impl(self) -> None:
+        return None
 
 
 class ApplicationCertificateTokenProvider(CloudInfoTokenProvider):
@@ -613,9 +612,8 @@ def _get_token_impl(self) -> Optional[dict]:
         token = self._msal_client.acquire_token_for_client(scopes=self._scopes)
         return self._valid_token_or_throw(token)
 
-    def _get_token_from_cache_impl(self) -> dict:
-        token = self._msal_client.acquire_token_silent(scopes=self._scopes, account=None)
-        return self._valid_token_or_none(token)
+    def _get_token_from_cache_impl(self) -> None:
+        return None
 
 
 class AzureIdentityTokenCredentialProvider(CloudInfoTokenProvider):

azure-kusto-data/azure/kusto/data/client_details.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from typing import List, Tuple, Optional
 
+from .env_utils import get_env
 from azure.kusto.data._version import VERSION
 
 NONE = "[none]"
@@ -23,8 +24,8 @@ def default_script() -> str:
 
 @functools.lru_cache(maxsize=1)
 def get_user_from_env() -> str:
-    user = os.environ.get("USERNAME", None)
-    domain = os.environ.get("USERDOMAIN", None)
+    user = get_env("USERNAME", optional=True)
+    domain = get_env("USERDOMAIN", optional=True)
     if domain and user:
         user = domain + "\\" + user
     if user:

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

@@ -0,0 +1,37 @@
+import os
+
+
+def get_env(*args, optional=False, default=None):
+    """Return the first environment variable that is defined."""
+    for arg in args:
+        if arg in os.environ:
+            return os.environ[arg]
+    if optional or default:
+        return default
+    raise ValueError("No environment variables found: {}".format(args))
+
+
+def set_env(key, value):
+    """Set the environment variable."""
+    os.environ[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
+    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
+    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
+    return result

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

@@ -1,12 +1,11 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License
-import os
-
 import pytest
 from azure.identity.aio import ClientSecretCredential as AsyncClientSecretCredential
 
 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 .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
 
@@ -162,8 +161,8 @@ async def test_msi_provider(self):
             print(" *** Skipped MSI Provider Test ***")
             return
 
-        user_msi_object_id = os.environ.get("MSI_OBJECT_ID")
-        user_msi_client_id = os.environ.get("MSI_CLIENT_ID")
+        user_msi_object_id = get_env("MSI_OBJECT_ID", optional=True)
+        user_msi_client_id = get_env("MSI_CLIENT_ID", optional=True)
 
         # system MSI
         with MsiTokenProvider(KUSTO_URI, is_async=True) as provider:
@@ -189,9 +188,9 @@ async def test_msi_provider(self):
     @aio_documented_by(TokenProviderTests.test_user_pass_provider)
     @pytest.mark.asyncio
     async def test_user_pass_provider(self):
-        username = os.environ.get("USER_NAME")
-        password = os.environ.get("USER_PASS")
-        auth = os.environ.get("USER_AUTH_ID", "organizations")
+        username = get_env("USER_NAME", optional=True)
+        password = get_env("USER_PASS", optional=True)
+        auth = get_env("USER_AUTH_ID", default="organizations")
 
         if username and password and auth:
             with UserPassTokenProvider(KUSTO_URI, auth, username, password, is_async=True) as provider:
@@ -228,18 +227,14 @@ def callback(x, x2, x3):
     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 = os.environ.get("APP_ID", "b699d721-4f6f-4320-bc9a-88d578dfe68f")
-        auth_id = os.environ.get("APP_AUTH_ID", "72f988bf-86f1-41af-91ab-2d7cd011db47")
-        app_key = os.environ.get("APP_KEY")
+        app_id = get_app_id(optional=True)
+        auth_id = get_auth_id(optional=True)
+        app_key = get_app_key(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:
                 token = await provider.get_token_async()
                 assert self.get_token_value(token) is not None
-
-                # Again through cache
-                token = provider._get_token_from_cache_impl()
-                assert self.get_token_value(token) is not None
         else:
             print(" *** Skipped App Id & Key Provider Test ***")
 
@@ -248,24 +243,20 @@ async def test_app_key_provider(self):
     async def test_app_cert_provider(self):
         # default details are for kusto-client-e2e-test-app
         # to invoke the test download the certs from Azure Portal
-        cert_app_id = os.environ.get("CERT_APP_ID", "b699d721-4f6f-4320-bc9a-88d578dfe68f")
-        cert_auth = os.environ.get("CERT_AUTH", "72f988bf-86f1-41af-91ab-2d7cd011db47")
-        thumbprint = os.environ.get("CERT_THUMBPRINT")
-        public_cert_path = os.environ.get("PUBLIC_CERT_PATH")
-        pem_key_path = os.environ.get("CERT_PEM_KEY_PATH")
+        cert_app_id = get_app_id(optional=True)
+        cert_auth = get_auth_id(optional=True)
+        thumbprint = get_env("CERT_THUMBPRINT", optional=True)
+        public_cert_path = get_env("CERT_PUBLIC_CERT_PATH", optional=True)
+        pem_key_path = get_env("CERT_PEM_KEY_PATH", optional=True)
 
-        if pem_key_path and thumbprint and cert_app_id:
+        if pem_key_path and thumbprint and cert_app_id and cert_auth:
             with open(pem_key_path, "rb") as file:
                 pem_key = file.read()
 
             with ApplicationCertificateTokenProvider(KUSTO_URI, cert_app_id, cert_auth, pem_key, thumbprint, is_async=True) as provider:
                 token = await provider.get_token_async()
                 assert self.get_token_value(token) is not None
 
-                # Again through cache
-                token = provider._get_token_from_cache_impl()
-                assert self.get_token_value(token) is not None
-
                 if public_cert_path:
                     with open(public_cert_path, "r") as file:
                         public_cert = file.read()
@@ -351,13 +342,9 @@ 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 = os.environ.get("APP_ID", "b699d721-4f6f-4320-bc9a-88d578dfe68f")
-        os.environ["AZURE_CLIENT_ID"] = app_id
-        auth_id = os.environ.get("APP_AUTH_ID", "72f988bf-86f1-41af-91ab-2d7cd011db47")
-        os.environ["AZURE_TENANT_ID"] = auth_id
-        app_key = os.environ.get("APP_KEY")
-        os.environ["AZURE_CLIENT_SECRET"] = app_key
-
+        app_id = get_app_id()
+        auth_id = get_auth_id()
+        app_key = get_app_key()
         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/e2e.py

@@ -12,6 +12,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 import KustoClient, KustoConnectionStringBuilder
 from azure.kusto.data._cloud_settings import CloudSettings, DEFAULT_DEV_KUSTO_SERVICE_RESOURCE_ID
 from azure.kusto.data._models import WellKnownDataSet
@@ -108,23 +109,17 @@ def engine_kcsb_from_env(cls, app_insights=False, is_async=False) -> KustoConnec
 
     @classmethod
     def setup_class(cls):
-        cls.engine_cs = os.environ.get("ENGINE_CONNECTION_STRING") or ""
-        cls.ai_engine_cs = os.environ.get("APPLICATION_INSIGHTS_ENGINE_CONNECTION_STRING") or ""
-        cls.app_id = os.environ.get("APP_ID")
-        if cls.app_id:
-            os.environ["AZURE_CLIENT_ID"] = cls.app_id
-        cls.app_key = os.environ.get("APP_KEY")
-        if cls.app_key:
-            os.environ["AZURE_CLIENT_SECRET"] = cls.app_key
-        cls.auth_id = os.environ.get("AUTH_ID")
-        if cls.auth_id:
-            os.environ["AZURE_TENANT_ID"] = cls.auth_id
-        os.environ["AZURE_AUTHORITY_HOST"] = "login.microsoftonline.com"
-        cls.test_db = os.environ.get("TEST_DATABASE")
-        cls.ai_test_db = os.environ.get("APPLICATION_INSIGHTS_TEST_DATABASE")  # name of e2e database could be changed
-
-        if not all([cls.engine_cs, cls.test_db, cls.ai_engine_cs, cls.ai_test_db]):
-            pytest.skip("E2E environment is missing")
+        cls.engine_cs = get_env("ENGINE_CONNECTION_STRING")
+        cls.ai_engine_cs = get_env("APPLICATION_INSIGHTS_ENGINE_CONNECTION_STRING")
+
+        cls.app_id = get_app_id()
+        cls.auth_id = get_auth_id()
+        cls.app_key = get_app_key()
+
+        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")  # name of e2e database could be changed
 
         # Init clients
         cls.streaming_test_table = "BigChunkus"