Skip to content

Commit

Permalink
Test connect with service credentials
Browse files Browse the repository at this point in the history
  • Loading branch information
steinitzu committed May 22, 2024
1 parent 763300b commit 1586d02
Showing 1 changed file with 38 additions and 5 deletions.
43 changes: 38 additions & 5 deletions tests/load/filesystem/test_azure_credentials.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand All @@ -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"
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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}")

0 comments on commit 1586d02

Please sign in to comment.