feat(elt-common): Pipeline testing harness

elt-common/src/elt_common/cli.py

@@ -7,7 +7,6 @@
 from typing import Literal, Optional, cast
 
 import click
-import requests
 
 from elt_common.pipeline import PipelinesProject
 from elt_common.pipeline_types import ELTIngestManifest
@@ -16,11 +15,6 @@
 
 LOGGER = logging.getLogger(__name__)
 
-# Disable ipv6 for anything using requests, which includes pyiceberg
-# This is done to solve a performance problem when using WSL
-# https://github.com/ISISNeutronMuon/analytics-data-platform/issues/433
-requests.packages.urllib3.util.connection.HAS_IPV6 = False
-
 
 @click.group(context_settings={"show_default": True})
 @click.option(

elt-common/src/elt_common/ingest.py

@@ -10,6 +10,7 @@
 
 import pyarrow as pa
 import pyarrow.compute as pc
+import requests
 from dotenv import load_dotenv, find_dotenv
 from pyiceberg.exceptions import NoSuchTableError
 
@@ -27,6 +28,12 @@
 LOGGER = logging.getLogger(__name__)
 
 
+# Disable ipv6 for anything using requests, which includes pyiceberg
+# This is done to solve a performance problem when using WSL
+# https://github.com/ISISNeutronMuon/analytics-data-platform/issues/433
+requests.packages.urllib3.util.connection.HAS_IPV6 = False
+
+
 def run_ingest(job: ELTIngestManifest) -> dict[str, int]:
     """Import the extract function, call it, and write results to Iceberg."""
 

elt-common/src/elt_common/pipeline.py

@@ -78,15 +78,15 @@ def _discover_jobs(warehouse_name: str, ingest_dir: Path):
     """
 
     return [
-        _create_ingest_manifest(warehouse_name, job_dir)
+        create_ingest_manifest(warehouse_name, job_dir)
         for domain_dir in ingest_dir.iterdir()
         if domain_dir.is_dir()
         for job_dir in domain_dir.iterdir()
         if job_dir.is_dir()
     ]
 
 
-def _create_ingest_manifest(warehouse_name: str, job_dir: Path) -> ELTIngestManifest:
+def create_ingest_manifest(warehouse_name: str, job_dir: Path) -> ELTIngestManifest:
     return ELTIngestManifest(
         warehouse_name=warehouse_name,
         name=job_dir.name,

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

@@ -2,18 +2,19 @@
 
 import tempfile
 import time
-from typing import Generator
 import urllib.parse
-import shutil
 import warnings
+from collections.abc import Generator
+from pathlib import Path
+from tempfile import TemporaryDirectory
 
-from minio import Minio
 import pytest
 import tenacity
+from minio import Minio
 
 from . import DEFAULT_RETRY_ARGS
 from .dlt import PyIcebergDestinationTestConfiguration
-from .lakekeeper import Settings, Server
+from .lakekeeper import Server, Settings
 from .sqlcatalog import SqlCatalogWarehouse
 
 
@@ -28,10 +29,11 @@ def warehouse(settings: Settings) -> Generator:
         )
 
     if settings.catalog_type == "sql":
-        warehouse = SqlCatalogWarehouse(settings.warehouse_name)
+        d = TemporaryDirectory()
+        warehouse = SqlCatalogWarehouse(settings.warehouse_name, Path(d.name))
 
         def cleanup_func():
-            shutil.rmtree(warehouse.workdir.name)
+            d.cleanup()
     else:
         server = Server(settings)
         storage_config = settings.storage_config()

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

@@ -0,0 +1,125 @@
+"""Pytest fixtures and utilities for e2e testing ingest pipelines.
+
+The `test_catalog` fixture creates a temporary file based iceberg warehouse,
+loads its catalog, and provides methods which tests can use to make assertions
+about the data which gets written to it.
+
+`run_test_ingest` exposes a method that runs the pipeline under test, optionally
+with configuration values provided as kwargs.
+
+Both rely on the test file being named 'test_<job>.py' and existing in the same
+directory as the pipeline it tests, in the elt-common directory structure:
+
+<warehouse_name>/
+|-- ingest/
+|   |-- <domain>/
+|   |   |-- <job>/
+|   |   |   |-- <job>.py
+|   |   |   |-- test_<job>.py
+"""
+
+import logging
+from pathlib import Path
+
+import pytest
+from pyiceberg.catalog import Catalog
+
+from elt_common.ingest import run_ingest
+from elt_common.pipeline import create_ingest_manifest
+from elt_common.testing.sqlcatalog import SqlCatalogWarehouse
+
+LOGGER = logging.getLogger(__name__)
+
+
+class AssertableCatalog:
+    """Wraps an iceberg catalog with convenience methods for making assertions
+    about the data in it"""
+
+    def __init__(self, catalog: Catalog):
+        self._catalog = catalog
+
+    def clean_catalog(self):
+        for ns in self._catalog.list_namespaces():
+            tables = self._catalog.list_tables(ns)
+            for qualified_table_name in tables:
+                self._catalog.purge_table(qualified_table_name)
+
+            self._catalog.drop_namespace(ns)
+
+    @property
+    def catalog(self):
+        return self._catalog
+
+    def assert_has_exact_tables(self, namespace: str, tables: list[str]):
+        actual = self._catalog.list_tables(namespace)
+        assert set(actual) == {(namespace, table) for table in tables}
+
+    def assert_has_columns(self, table_id: tuple[str, ...], column_names):
+        assert self._catalog.table_exists(table_id), f"{table_id} doesn't exist"
+        t = self._catalog.load_table(table_id)
+        for c in column_names:
+            assert c in t.schema().column_names
+
+    def assert_has_n_rows(self, table_id: tuple[str, ...], n: int):
+        assert self.get_num_rows(table_id) == n
+
+    def get_num_rows(self, table_id: tuple[str, ...]):
+        assert self._catalog.table_exists(table_id)
+        t = self._catalog.load_table(table_id)
+        return t.scan().count()
+
+
+@pytest.fixture(scope="session")
+def sql_warehouses(request, tmp_path_factory):
+    test_dir = tmp_path_factory.mktemp("warehouses")
+    test_paths = (t.path for t in request.session.items)
+    test_warehouse_names = {_get_warehouse_name_from_test_filepath(tp) for tp in test_paths}
+    LOGGER.debug(f"Creating {test_warehouse_names} warehouses in {test_dir}")
+
+    warehouses = {
+        warehouse_name: SqlCatalogWarehouse(warehouse_name, test_dir)
+        for warehouse_name in test_warehouse_names
+    }
+
+    return warehouses
+
+
+@pytest.fixture
+def test_catalog(sql_warehouses, request, monkeypatch):
+    test_warehouse_name = _get_warehouse_name_from_test_filepath(request.path)
+    warehouse = sql_warehouses[test_warehouse_name]
+
+    monkeypatch.setenv("PYICEBERG_CATALOG__DEFAULT__TYPE", "sql")
+    monkeypatch.setenv("PYICEBERG_CATALOG__DEFAULT__URI", warehouse.uri)
+    monkeypatch.setenv("PYICEBERG_CATALOG__DEFAULT__WAREHOUSE", test_warehouse_name)
+
+    catalog = warehouse.connect()
+    ret = AssertableCatalog(catalog)
+
+    try:
+        yield ret
+    finally:
+        try:
+            ret.clean_catalog()
+        finally:
+            catalog.close()
+
+
+@pytest.fixture
+def run_test_ingest(request, monkeypatch):
+    w = _get_warehouse_name_from_test_filepath(request.path)
+    manifest = create_ingest_manifest(w, request.path.parent)
+
+    def run(**config_vars):
+        LOGGER.debug(f"Running ingest for manifest {manifest} with config vars {config_vars}")
+        for k, v in config_vars.items():
+            monkeypatch.setenv(f"{manifest.name}__{k}", v)
+
+        run_ingest(manifest)
+
+    return run
+
+
+def _get_warehouse_name_from_test_filepath(fp: Path):
+    """Relies on the test file being '<warehouse>/ingest/<domain>/<job>/test_<job>.py'"""
+    return fp.parent.parent.parent.parent.name

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

@@ -1,16 +1,17 @@
-from tempfile import TemporaryDirectory
+from pathlib import Path
 
-from elt_common.dlt_destinations.pyiceberg.configuration import PyIcebergSqlCatalogCredentials
 from pyiceberg.catalog import Catalog as PyIcebergCatalog
 from pyiceberg.catalog import load_catalog
 
+from elt_common.dlt_destinations.pyiceberg.configuration import PyIcebergSqlCatalogCredentials
+
 
 class SqlCatalogWarehouse:
-    def __init__(self, warehouse_name: str):
+    def __init__(self, warehouse_name: str, workdir: Path):
         self.name = warehouse_name
-        self.workdir = TemporaryDirectory()
-        self.uri = f"sqlite:///{self.workdir.name}/{warehouse_name}.db"
-        self.warehouse_path = f"file://{self.workdir.name}/{warehouse_name}"
+        self.workdir = workdir
+        self.uri = f"sqlite:///{self.workdir}/{warehouse_name}.db"
+        self.warehouse_path = f"file://{self.workdir}/{warehouse_name}"
 
     def connect(self) -> PyIcebergCatalog:
         """Connect to the warehouse in the catalog"""

elt-pipelines/README.md

@@ -95,6 +95,17 @@ class Extract(BaseExtract):
 There is functionality built in to `elt_common` for some common types of pipeline. These can be found in the
 [`sources` package](../elt-common/src/elt_common/sources).
 
+### Testing
+
+To test a pipeline, create a `test_<job name>.py` file in the same directory as the pipeline script and use
+the `test_catalog` and `run_test_ingest` fixtures defined in `elt_common.testing.pipelines`.
+
+Tests are invoked with `pytest --disable-plugin-autoload -p elt_common.testing.pipelines`. The plugin arguments
+are needed temporarily, until the `elt_common.testing.fixtures` module (which provides fixtures for `dlt` based
+testing) is removed.
+
+**Configuration values need to be set for any pipeline(s) under test, as per [above](#configuration).**
+
 ## Directory structure
 
 The project uses the following directory structure:

elt-pipelines/facility_ops/ingest/accelerator/statusdisplay/test_statusdisplay.py

@@ -0,0 +1,42 @@
+from elt_common.testing.pipelines import AssertableCatalog
+
+_namespace = "accelerator_statusdisplay"
+_table_name = "elt_cycles"
+_table_id = (_namespace, _table_name)
+
+
+def test_expected_columns_created(test_catalog: AssertableCatalog, run_test_ingest):
+    run_test_ingest()
+    test_catalog.assert_has_columns(
+        _table_id,
+        [
+            "id",
+            "label",
+            "status",
+            "phases",
+            "phases.element",  # phases is a list type
+        ],
+    )
+    num_rows = test_catalog.get_num_rows(_table_id)
+    assert num_rows >= 153, f"Found {num_rows} cycles, expected at least 153"
+
+    test_catalog.clean_catalog()
+
+    assert not test_catalog.catalog.table_exists(_table_id), (
+        "Table should have been cleaned up"
+    )
+
+
+def test_multiple_runs(test_catalog: AssertableCatalog, run_test_ingest):
+    run_test_ingest()
+    n = test_catalog.get_num_rows(_table_id)
+    run_test_ingest()
+    run_test_ingest()
+
+    # Check runs are overwriting, not appending
+    test_catalog.assert_has_n_rows(_table_id, n)
+
+    nss = test_catalog.catalog.list_namespaces()
+    assert len(nss) == 1, "There should be a single namespace"
+    tables = test_catalog.catalog.list_tables(_namespace)
+    assert len(tables) == 1, "There should be a single table"

elt-pipelines/fase/ingest/fase/proposal/test_proposal.py

@@ -0,0 +1,31 @@
+import json
+import pytest
+
+from elt_common.testing.pipelines import AssertableCatalog
+
+_ns = "fase_proposal"
+
+
+def test_expected_columns_created(test_catalog: AssertableCatalog, run_test_ingest):
+    run_test_ingest(tables=json.dumps(["call"]))
+    expected_columns = [
+        "call_id",
+        "call_short_code",
+        "start_call",
+        "end_call",
+        "start_review",
+        "end_review",
+    ]
+    test_catalog.assert_has_columns((_ns, "call"), expected_columns)
+
+
+def test_multiple_tables(test_catalog: AssertableCatalog, run_test_ingest):
+    tables = ["call", "countries", "questions"]
+    run_test_ingest(tables=json.dumps(tables), row_limit="5")
+    test_catalog.assert_has_exact_tables(_ns, tables)
+
+
+@pytest.mark.parametrize("row_limit", [1, 5, 20])
+def test_row_limit_applied(test_catalog: AssertableCatalog, run_test_ingest, row_limit):
+    run_test_ingest(row_limit=str(row_limit), tables=json.dumps(["call"]))
+    test_catalog.assert_has_n_rows((_ns, "call"), row_limit)

elt-pipelines/pyproject.toml

@@ -41,4 +41,5 @@ elt-common = { path = "../elt-common", editable = true }
 [dependency-groups]
 dev = [
     "prek>=0.4.5",
+    "pytest>=9.1.1",
 ]