Implement testing for Opralog extract_and_load script

Author: martyngigg

elt-common/pyproject.toml

@@ -10,28 +10,28 @@ readme = "README.md"
 requires-python = ">=3.13"
 dependencies = ["dlt[parquet,s3]~=1.15.0", "pyiceberg~=0.9.1"]
 
+
 [project.optional-dependencies]
 iceberg-maintenance = ["click~=8.2.1", "sqlalchemy~=2.0.43", "trino~=0.336.0"]
-m365 = [
-    "authlib~=1.6.1",
-    "httpx~=0.28.1",
-    "tenacity~=9.1.2",
-]
+m365 = ["authlib~=1.6.1", "httpx~=0.28.1", "tenacity~=9.1.2"]
+
+[project.entry-points.pytest11]
+elt_common = "elt_common.testing.fixtures"
 
 [project.scripts]
 iceberg-maintenance = "elt_common.iceberg.maintenance:cli"
 
 [dependency-groups]
 dev = [
-    "minio~=7.2.16",
-    "prek>=0.1.6",
-    "pydantic-settings~=2.10.1",
-    "pytest~=8.3.5",
-    "pytest-httpx~=0.35.0",
-    "pytest-mock~=3.14.1",
-    "requests~=2.32.3",
-    "requests-mock~=1.12.1",
-    "ruff~=0.11.11",
+  "minio~=7.2.16",
+  "prek>=0.1.6",
+  "pydantic-settings~=2.10.1",
+  "pytest~=8.3.5",
+  "pytest-httpx~=0.35.0",
+  "pytest-mock~=3.14.1",
+  "requests~=2.32.3",
+  "requests-mock~=1.12.1",
+  "ruff~=0.11.11",
 ]
 
 

elt-common/src/elt_common/dlt_destinations/pyiceberg/__init__.py

@@ -0,0 +1,14 @@
+from contextlib import contextmanager
+from typing import cast
+
+import dlt
+
+from .pyiceberg import PyIcebergClient
+
+
+@contextmanager
+def iceberg_catalog(pipeline: dlt.Pipeline):
+    """Given a pipeline return the Iceberg catalog client for this pipeline"""
+    with pipeline.destination_client() as client:
+        assert isinstance(client, PyIcebergClient)
+        yield cast(PyIcebergClient, client).iceberg_catalog

elt-common/src/elt_common/testing/__init__.py

@@ -0,0 +1,96 @@
+import dataclasses
+import urllib.parse
+from typing import Any, Dict
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+import tenacity
+
+DEFAULT_RETRY_ARGS = {
+    "wait": tenacity.wait_exponential(max=10),
+    "stop": tenacity.stop_after_attempt(5),
+    "reraise": True,
+}
+
+
+@dataclasses.dataclass
+class Endpoint:
+    raw_value: str
+    internal_netloc: str
+
+    def __add__(self, path: str) -> "Endpoint":
+        return Endpoint(self.raw_value + path, self.internal_netloc)
+
+    def __str__(self) -> str:
+        return self.raw_value
+
+    def value(self, *, use_internal_netloc: bool) -> str:
+        if use_internal_netloc:
+            fragments = list(urllib.parse.urlparse(self.raw_value))
+            fragments[1] = self.internal_netloc
+            return urllib.parse.urlunparse(fragments)
+        else:
+            return self.raw_value
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="tests_",
+    )
+
+    # iceberg catalog
+    # The default values assume the docker-compose.yml in the infra/local has been used.
+    # These are provided for the convenience of easily running a debugger without having
+    # to set up remote debugging
+    host_netloc: str = "localhost:50080"
+    docker_netloc: str = "adp-router:50080"
+    s3_access_key: str = "adpsuperuser"
+    s3_secret_key: str = "adppassword"
+    s3_bucket: str = "e2e-tests"
+    s3_endpoint: str = "http://adp-router:59000"
+    s3_region: str = "local-01"
+    s3_path_style_access: bool = True
+    openid_client_id: str = "machine-infra"
+    openid_client_secret: str = "s3cr3t"
+    openid_scope: str = "lakekeeper"
+    project_id: str = "c4fcd44f-7ce7-4446-9f7c-dcc7ba76dd22"
+    warehouse_name: str = "e2e_tests"
+
+    # trino
+    trino_http_scheme: str = "https"
+    trino_host: str = "localhost"
+    trino_port: str = "58443"
+    trino_user: str = "machine-infra"
+    trino_password: str = "s3cr3t"
+
+    @property
+    def lakekeeper_url(self) -> Endpoint:
+        return Endpoint(f"http://{self.host_netloc}/iceberg", self.docker_netloc)
+
+    @property
+    def openid_provider_uri(self) -> Endpoint:
+        return Endpoint(
+            f"http://{self.host_netloc}/auth/realms/analytics-data-platform", self.docker_netloc
+        )
+
+    def storage_config(self) -> Dict[str, Any]:
+        return {
+            "warehouse-name": self.warehouse_name,
+            "storage-credential": {
+                "type": "s3",
+                "credential-type": "access-key",
+                "aws-access-key-id": self.s3_access_key,
+                "aws-secret-access-key": self.s3_secret_key,
+            },
+            "storage-profile": {
+                "type": "s3",
+                "bucket": self.s3_bucket,
+                "key-prefix": "",
+                "assume-role-arn": "",
+                "endpoint": self.s3_endpoint,
+                "region": self.s3_region,
+                "path-style-access": self.s3_path_style_access,
+                "flavor": "s3-compat",
+                "sts-enabled": False,
+            },
+            "delete-profile": {"type": "hard"},
+        }

elt-common/src/elt_common/testing/dlt.py

@@ -0,0 +1,101 @@
+from collections.abc import MutableMapping
+from dataclasses import dataclass
+import os
+from typing import List
+
+import dlt
+from dlt.common.configuration.providers import (
+    ConfigProvider,
+    EnvironProvider,
+)
+from dlt.common.destination import (
+    TDestinationReferenceArg,
+)
+from dlt.common.runtime.run_context import RunContext
+
+from .lakekeeper import Warehouse
+
+
+def configure_dlt_for_testing():
+    """Set up dlt for a testing environment"""
+
+    def initial_providers(self) -> List[ConfigProvider]:
+        # Use only environmental variables to configure
+        return [
+            EnvironProvider(),
+        ]
+
+    os.environ["RUNTIME__DLTHUB_TELEMETRY"] = "false"
+    RunContext.initial_providers = initial_providers  # type: ignore[method-assign]
+
+
+@dataclass
+class PyIcebergDestinationTestConfiguration:
+    """Class for defining test setup for pyiceberg destination."""
+
+    warehouse: Warehouse
+    destination: TDestinationReferenceArg = "elt_common.dlt_destinations.pyiceberg"
+
+    def setup(self, environ: MutableMapping = os.environ) -> None:
+        """Sets up environment variables for this destination configuration
+
+        Defaults to to os.environ
+        """
+        server, server_settings = self.warehouse.server, self.warehouse.server.settings
+        environ["DESTINATION__PYICEBERG__CREDENTIALS__URI"] = str(
+            self.warehouse.server.catalog_endpoint()
+        )
+        environ["DESTINATION__PYICEBERG__CREDENTIALS__PROJECT_ID"] = str(
+            self.warehouse.server.settings.project_id
+        )
+        environ.setdefault("DESTINATION__PYICEBERG__CREDENTIALS__WAREHOUSE", self.warehouse.name)
+        environ.setdefault(
+            "DESTINATION__PYICEBERG__CREDENTIALS__OAUTH2_SERVER_URI",
+            str(server.token_endpoint),
+        )
+        environ.setdefault(
+            "DESTINATION__PYICEBERG__CREDENTIALS__CLIENT_ID",
+            server_settings.openid_client_id,
+        )
+        environ.setdefault(
+            "DESTINATION__PYICEBERG__CREDENTIALS__CLIENT_SECRET",
+            server_settings.openid_client_secret,
+        )
+        environ.setdefault(
+            "DESTINATION__PYICEBERG__CREDENTIALS__SCOPE",
+            server_settings.openid_scope,
+        )
+        environ.setdefault("DESTINATION__PYICEBERG__BUCKET_URL", self.warehouse.bucket_url)
+        # Avoid collisons on table names when the same ones are created/deleted in quick
+        # succession
+        environ.setdefault(
+            "DESTINATION__PYICEBERG__TABLE_LOCATION_LAYOUT",
+            "{location_tag}/{dataset_name}/{table_name}",
+        )
+
+    def setup_pipeline(
+        self,
+        pipeline_name: str,
+        dataset_name: str = None,
+        dev_mode: bool = False,
+        **kwargs,
+    ) -> dlt.Pipeline:
+        """Convenience method to setup pipeline with this configuration"""
+        self.setup()
+        return dlt.pipeline(
+            pipeline_name=pipeline_name,
+            pipelines_dir=kwargs.pop("pipelines_dir", None),
+            destination=self.destination,
+            dataset_name=(dataset_name if dataset_name is not None else pipeline_name + "_data"),
+            dev_mode=dev_mode,
+            **kwargs,
+        )
+
+    def clean_catalog(self):
+        """Clean the destination catalog of all namespaces and tables"""
+        catalog = self.warehouse.connect()
+        for ns_name in catalog.list_namespaces():
+            tables = catalog.list_tables(ns_name)
+            for qualified_table_name in tables:
+                catalog.purge_table(qualified_table_name)
+            catalog.drop_namespace(ns_name)

elt-common/src/elt_common/testing/fixtures.py

@@ -0,0 +1,108 @@
+"""A collection of utilities to support testing against this library"""
+
+import tempfile
+import time
+from typing import Generator
+import urllib.parse
+import warnings
+
+from minio import Minio
+import pytest
+import requests
+import tenacity
+
+from . import DEFAULT_RETRY_ARGS
+from .dlt import PyIcebergDestinationTestConfiguration
+from .lakekeeper import Settings, Server, Warehouse
+
+
+@pytest.fixture(scope="session")
+def settings() -> Settings:
+    return Settings()
+
+
+@pytest.fixture(scope="session")
+def token_endpoint(settings: Settings) -> str:
+    response = requests.get(str(settings.openid_provider_uri + "/.well-known/openid-configuration"))
+    response.raise_for_status()
+    return response.json()["token_endpoint"]
+
+
+@pytest.fixture(scope="session")
+def access_token(settings: Settings, token_endpoint: str) -> str:
+    response = requests.post(
+        token_endpoint,
+        data={
+            "grant_type": "client_credentials",
+            "client_id": settings.openid_client_id,
+            "client_secret": settings.openid_client_secret,
+            "scope": settings.openid_scope,
+        },
+    )
+    response.raise_for_status()
+    return response.json()["access_token"]
+
+
+@pytest.fixture(scope="session")
+def server(settings: Settings, access_token: str) -> Server:
+    return Server(access_token, settings)
+
+
+@pytest.fixture(scope="session")
+def warehouse(settings: Settings, server: Server) -> Generator:
+    if not settings.warehouse_name:
+        raise ValueError("Empty 'warehouse_name' is not allowed.")
+
+    storage_config = settings.storage_config()
+    # Ensure bucket exists
+    s3_hostname = urllib.parse.urlparse(storage_config["storage-profile"]["endpoint"]).netloc
+    minio_client = Minio(
+        endpoint=s3_hostname,
+        access_key=storage_config["storage-credential"]["aws-access-key-id"],
+        secret_key=storage_config["storage-credential"]["aws-secret-access-key"],
+        secure=False,
+    )
+    bucket_name = storage_config["storage-profile"]["bucket"]
+    if not minio_client.bucket_exists(bucket_name=bucket_name):
+        minio_client.make_bucket(bucket_name=bucket_name)
+        print(f"Bucket {bucket_name} created.")
+
+    warehouse = server.create_warehouse(
+        settings.warehouse_name, server.settings.project_id, storage_config
+    )
+    print(f"Warehouse {warehouse.project_id} created.")
+    try:
+        yield warehouse
+    finally:
+
+        @tenacity.retry(**DEFAULT_RETRY_ARGS)
+        def _remove_bucket(bucket_name):
+            minio_client.remove_bucket(bucket_name=bucket_name)
+
+        try:
+            # Allow a brief pause for the test operations to complete
+            time.sleep(1)
+            server.purge_warehouse(warehouse)
+            server.delete_warehouse(warehouse)
+            _remove_bucket(bucket_name)
+
+        except RuntimeError as exc:
+            warnings.warn(
+                f"Error deleting test warehouse '{str(warehouse.project_id)}'. It may need to be removed manually."
+            )
+            warnings.warn(f"Error:\n{str(exc)}")
+
+
+@pytest.fixture
+def destination_config(warehouse: Warehouse):
+    destination_config = PyIcebergDestinationTestConfiguration(warehouse)
+    try:
+        yield destination_config
+    finally:
+        destination_config.clean_catalog()
+
+
+@pytest.fixture
+def pipelines_dir():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        yield tmp_dir