From 1586d023383f39a670c7c089a5bea453ce80cc29 Mon Sep 17 00:00:00 2001
From: Steinthor Palsson <steini90@gmail.com>
Date: Wed, 22 May 2024 18:58:57 -0400
Subject: [PATCH] Test connect with service credentials

---
 .../load/filesystem/test_azure_credentials.py | 43 ++++++++++++++++---
 1 file changed, 38 insertions(+), 5 deletions(-)

diff --git a/tests/load/filesystem/test_azure_credentials.py b/tests/load/filesystem/test_azure_credentials.py
index 885c6aba4f..ff1288f7f3 100644
--- a/tests/load/filesystem/test_azure_credentials.py
+++ b/tests/load/filesystem/test_azure_credentials.py
@@ -1,9 +1,10 @@
-from typing import Dict
+from typing import Dict, Optional
 from urllib.parse import parse_qs
 from uuid import uuid4
 
 import pytest
 
+import dlt
 from dlt.common import pendulum
 from dlt.common.time import ensure_pendulum_datetime
 from dlt.common.configuration import resolve_configuration, ConfigFieldMissingException
@@ -13,7 +14,8 @@
     AzureServicePrincipalCredentialsWithoutDefaults,
     AzureCredentialsWithoutDefaults,
 )
-from tests.load.utils import ALL_FILESYSTEM_DRIVERS
+from dlt.common.storages.configuration import FilesystemConfiguration
+from tests.load.utils import ALL_FILESYSTEM_DRIVERS, AZ_BUCKET
 from tests.common.configuration.utils import environment
 from tests.utils import preserve_environ, autouse_test_storage
 from dlt.common.storages.fsspec_filesystem import fsspec_from_config
@@ -25,6 +27,25 @@
     pytest.skip("az filesystem driver not configured", allow_module_level=True)
 
 
+@pytest.fixture
+def az_service_principal_config() -> Optional[FilesystemConfiguration]:
+    """FS config with alternate azure credentials format if available in environment"""
+    credentials = AzureServicePrincipalCredentialsWithoutDefaults(
+        azure_tenant_id=dlt.config.get("tests.az_sp_tenant_id", str),
+        azure_client_id=dlt.config.get("tests.az_sp_client_id", str),
+        azure_client_secret=dlt.config.get("tests.az_sp_client_secret", str),
+        azure_storage_account_name=dlt.config.get("tests.az_sp_storage_account_name", str),
+    )
+    try:
+        credentials = resolve_configuration(credentials)
+    except ConfigFieldMissingException:
+        pytest.skip("Azure service principal credentials not available in environment")
+        return None
+    cfg = FilesystemConfiguration(bucket_url=AZ_BUCKET, credentials=credentials)
+
+    return resolve_configuration(cfg)
+
+
 def test_azure_credentials_from_account_key(environment: Dict[str, str]) -> None:
     environment["CREDENTIALS__AZURE_STORAGE_ACCOUNT_NAME"] = "fake_account_name"
     environment["CREDENTIALS__AZURE_STORAGE_ACCOUNT_KEY"] = "QWERTYUIOPASDFGHJKLZXCVBNM1234567890"
@@ -124,9 +145,6 @@ def test_azure_service_principal_credentials(environment: Dict[str, str]) -> Non
     }
 
 
-from dlt.common.storages.configuration import FilesystemConfiguration
-
-
 def test_azure_filesystem_configuration_service_principal(environment: Dict[str, str]) -> None:
     """Filesystem config resolves correct credentials type"""
     environment["CREDENTIALS__AZURE_STORAGE_ACCOUNT_NAME"] = "fake_account_name"
@@ -163,3 +181,18 @@ def test_azure_filesystem_configuration_sas_token(environment: Dict[str, str]) -
 
     assert fs.sas_token == "?" + environment["CREDENTIALS__AZURE_STORAGE_SAS_TOKEN"]
     assert fs.account_name == environment["CREDENTIALS__AZURE_STORAGE_ACCOUNT_NAME"]
+
+
+def test_azure_service_principal_fs_operations(
+    az_service_principal_config: Optional[FilesystemConfiguration],
+) -> None:
+    """Test connecting to azure filesystem with service principal credentials"""
+    config = az_service_principal_config
+    fs, bucket = fsspec_from_config(config)
+
+    fn = uuid4().hex
+    # Try some file ops to see if the credentials work
+    fs.touch(f"{bucket}/{fn}")
+    files = fs.ls(bucket)
+    assert f"{bucket}/{fn}" in files
+    fs.delete(f"{bucket}/{fn}")