stack/tests/aco/test_lake_deploy_exercise.py

"""Exercise aco.lake.deploy — deploy_schemas, _ensure_namespace, _create_iceberg_table."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

from aco.lake.deploy import (
    _ensure_namespace,
    deploy_schemas,
)


class TestEnsureNamespace:
    def test_creates(self):
        cat = MagicMock()
        _ensure_namespace(cat, "raw")
        cat._get_iceberg_catalog.return_value.create_namespace.assert_called()

    def test_already_exists(self):
        cat = MagicMock()
        cat._get_iceberg_catalog.return_value.create_namespace.side_effect = Exception(
            "exists"
        )
        _ensure_namespace(cat, "raw")  # no error


class TestDeployWithConfig:
    @patch("aco.lake.catalog.Catalog")
    @patch("conf.cfg")
    def test_nessie_config(self, mc_cfg, mc_catalog):
        mc_cfg.lake.nessie.catalog_uri = "http://nessie:19120/api/v1"
        mc_cfg.lake.warehouse = "s3://bucket"
        cat = mc_catalog.return_value
        cat.schemas.return_value = []

        result = deploy_schemas(catalog_type="nessie", dry_run=True)
        assert result == {}

    @patch("aco.lake.catalog.Catalog")
    @patch("conf.cfg")
    def test_polaris_config(self, mc_cfg, mc_catalog):
        mc_cfg.lake.polaris.catalog_uri = "http://polaris:8181/api/v1"
        mc_cfg.lake.warehouse = "s3://bucket"
        cat = mc_catalog.return_value
        cat.schemas.return_value = []

        result = deploy_schemas(catalog_type="polaris", dry_run=True)
        assert result == {}


class TestCreateIcebergTable:
    def test_creates(self):
        import sys

        mock_pyiceberg = MagicMock()
        with patch.dict(
            sys.modules,
            {
                "pyiceberg": mock_pyiceberg,
                "pyiceberg.schema": mock_pyiceberg.schema,
                "pyiceberg.types": mock_pyiceberg.types,
            },
        ):
            from aco.lake.deploy import _create_iceberg_table

            cat = MagicMock()
            model = MagicMock()
            field_info = MagicMock()
            field_info.annotation = str
            model.model_fields = {"col1": field_info, "col2": field_info}
            with patch("aco.lake.deploy._resolve_iceberg_type", return_value="string"):
                _create_iceberg_table(cat, "ns.schema.table", model)
            cat._get_iceberg_catalog.return_value.create_table.assert_called_once()


class TestDeployTables:
    @patch("aco.lake.catalog.Catalog")
    @patch("conf.cfg")
    def test_dry_run(self, mc_cfg, mc_catalog):
        cat = mc_catalog.return_value
        cat.schemas.return_value = ["raw"]
        cat.tables.return_value = ["ns.raw.table1"]
        model = MagicMock()
        model.model_fields = {"col1": MagicMock()}
        cat.model.return_value = model

        result = deploy_schemas(
            catalog_uri="http://test:8181/api/v1",
            warehouse="wh",
            dry_run=True,
        )
        assert "raw" in result
        assert "ns.raw.table1" in result["raw"]

    @patch("aco.lake.catalog.Catalog")
    @patch("aco.lake.deploy._ensure_namespace")
    @patch("aco.lake.deploy._create_iceberg_table")
    @patch("conf.cfg")
    def test_creates(self, mc_cfg, mc_create, mc_ns, mc_catalog):
        cat = mc_catalog.return_value
        cat.schemas.return_value = ["raw"]
        cat.tables.return_value = ["ns.raw.table1"]
        model = MagicMock()
        cat.model.return_value = model

        result = deploy_schemas(
            catalog_uri="http://test:8181/api/v1",
            warehouse="wh",
            dry_run=False,
        )
        assert "raw" in result
        mc_create.assert_called_once()

    @patch("aco.lake.catalog.Catalog")
    @patch("aco.lake.deploy._ensure_namespace")
    @patch("aco.lake.deploy._create_iceberg_table")
    @patch("conf.cfg")
    def test_already_exists(self, mc_cfg, mc_create, mc_ns, mc_catalog):
        cat = mc_catalog.return_value
        cat.schemas.return_value = ["raw"]
        cat.tables.return_value = ["ns.raw.table1"]
        model = MagicMock()
        cat.model.return_value = model
        mc_create.side_effect = Exception("already exists")

        result = deploy_schemas(
            catalog_uri="http://test:8181/api/v1",
            warehouse="wh",
            dry_run=False,
        )
        assert "ns.raw.table1" in result["raw"]

    @patch("aco.lake.catalog.Catalog")
    @patch("aco.lake.deploy._ensure_namespace")
    @patch("aco.lake.deploy._create_iceberg_table")
    @patch("conf.cfg")
    def test_create_failure(self, mc_cfg, mc_create, mc_ns, mc_catalog):
        cat = mc_catalog.return_value
        cat.schemas.return_value = ["raw"]
        cat.tables.return_value = ["ns.raw.table1"]
        model = MagicMock()
        cat.model.return_value = model
        mc_create.side_effect = Exception("real error")

        result = deploy_schemas(
            catalog_uri="http://test:8181/api/v1",
            warehouse="wh",
            dry_run=False,
        )
        assert result["raw"] == []

    @patch("aco.lake.catalog.Catalog")
    @patch("conf.cfg")
    def test_empty_schema(self, mc_cfg, mc_catalog):
        cat = mc_catalog.return_value
        cat.schemas.return_value = ["empty_schema"]
        cat.tables.return_value = []

        result = deploy_schemas(
            catalog_uri="http://test:8181/api/v1",
            warehouse="wh",
            dry_run=True,
        )
        assert result == {}