stack/tests/test_dab_generator.py

"""Tests for the Databricks Asset Bundle generator.

Covers: YAML structure, Python resource loader, pipeline metadata,
dependency ordering, schema filtering, idempotency, and targets.
"""

from __future__ import annotations

import sys
from pathlib import Path

import yaml

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "dev" / "scripts"))


class TestPipelineMetadata:
    """Pipeline.upstream and cluster_profile."""

    def test_all_pipelines_have_upstream(self):
        from aco.pipe import registry

        for name, pipe in registry.items():
            assert hasattr(pipe, "upstream"), f"{name} missing upstream"
            assert isinstance(pipe.upstream, list), f"{name}.upstream not a list"

    def test_all_pipelines_have_cluster_profile(self):
        from aco.pipe import registry

        for name, pipe in registry.items():
            assert hasattr(pipe, "cluster_profile"), f"{name} missing cluster_profile"

    def test_no_cycles_in_upstream(self):
        from aco.pipe import registry

        visited: set[str] = set()
        path: set[str] = set()

        def visit(name: str) -> None:
            if name in path:
                raise AssertionError(f"Cycle detected: {name} in {path}")
            if name in visited:
                return
            path.add(name)
            pipe = registry.get(name)
            if pipe:
                for dep in pipe.upstream:
                    visit(dep)
            path.remove(name)
            visited.add(name)

        for name in registry:
            visit(name)

    def test_upstream_refs_exist_in_registry(self):
        from aco.pipe import registry

        for name, pipe in registry.items():
            for dep in pipe.upstream:
                assert dep in registry, (
                    f"{name} references upstream '{dep}' not in registry"
                )

    def test_root_pipelines_have_empty_upstream(self):
        from aco.pipe import registry

        assert registry["input_layer"].upstream == []


class TestDabYaml:
    """Generated databricks.yml has correct structure."""

    def _yaml(self) -> dict:
        from backends.databricks import emit

        from conf import cfg

        files = emit(cfg._data)
        return yaml.safe_load(files["databricks.yml"])

    def test_generates_valid_yaml(self):
        data = self._yaml()
        assert isinstance(data, dict)

    def test_has_required_top_level_keys(self):
        data = self._yaml()
        for key in ("bundle", "python", "workspace", "targets"):
            assert key in data, f"Missing top-level key: {key}"

    def test_bundle_name(self):
        assert self._yaml()["bundle"]["name"] == "stack"

    def test_python_resources_entry_point(self):
        data = self._yaml()
        assert "resources" in data["python"]
        assert "resources:load_resources" in data["python"]["resources"]

    def test_no_inline_resources(self):
        """Resources come from Python, not inline YAML."""
        data = self._yaml()
        assert "resources" not in data or data.get("resources") is None

    def test_no_variables_in_yaml(self):
        """Variables defined in Python resource loader, not YAML."""
        data = self._yaml()
        assert "variables" not in data


class TestPythonResources:
    """Python resource loader produces correct typed resources."""

    def _resources(self):
        from bundle.resources import load_resources

        return load_resources()

    def test_load_resources_returns_resources(self):
        from databricks.bundles.core import Resources

        r = self._resources()
        assert isinstance(r, Resources)

    def test_has_job(self):
        r = self._resources()
        assert "stack_pipelines" in r.jobs

    def test_job_has_all_pipeline_tasks(self):
        from aco.pipe import registry

        r = self._resources()
        job = r.jobs["stack_pipelines"]
        task_keys = {t.task_key for t in job.tasks}
        for name in registry:
            assert name in task_keys, f"Pipeline '{name}' missing from tasks"

    def test_task_count_matches_registry(self):
        from aco.pipe import registry

        r = self._resources()
        assert len(r.jobs["stack_pipelines"].tasks) == len(registry)

    def test_tasks_use_python_wheel_task(self):
        r = self._resources()
        for task in r.jobs["stack_pipelines"].tasks:
            assert task.python_wheel_task is not None, (
                f"Task {task.task_key} missing python_wheel_task"
            )
            assert task.python_wheel_task.package_name == "stack"
            assert task.python_wheel_task.entry_point == "cli"

    def test_tasks_have_libraries(self):
        r = self._resources()
        for task in r.jobs["stack_pipelines"].tasks:
            assert task.libraries, f"Task {task.task_key} missing libraries"

    def test_dependencies_match_upstream(self):
        from aco.pipe import registry

        r = self._resources()
        task_map = {t.task_key: t for t in r.jobs["stack_pipelines"].tasks}
        for name, pipe in registry.items():
            task = task_map[name]
            if pipe.upstream:
                deps = sorted(d.task_key for d in task.depends_on)
                assert deps == sorted(pipe.upstream)
            else:
                assert not task.depends_on

    def test_schemas_count(self):
        r = self._resources()
        assert len(r.schemas) == 14, f"Expected 14 schemas, got {len(r.schemas)}"

    def test_schemas_include_core(self):
        r = self._resources()
        assert "core" in r.schemas

    def test_no_legacy_schemas(self):
        r = self._resources()
        for legacy in ("alr", "ccsr", "ccw", "pfs", "ssp", "reach"):
            assert legacy not in r.schemas, (
                f"Legacy schema '{legacy}' should not be deployed"
            )

    def test_has_staging_volume(self):
        r = self._resources()
        assert "staging" in r.volumes

    def test_job_has_schedule(self):
        r = self._resources()
        job = r.jobs["stack_pipelines"]
        assert job.schedule is not None
        assert job.schedule.quartz_cron_expression == "0 0 6 * * ?"


class TestTargets:
    """Target isolation."""

    def _yaml(self) -> dict:
        from backends.databricks import emit

        from conf import cfg

        return yaml.safe_load(emit(cfg._data)["databricks.yml"])

    def test_three_targets(self):
        assert set(self._yaml()["targets"].keys()) == {"dev", "staging", "prod"}

    def test_dev_is_default(self):
        assert self._yaml()["targets"]["dev"].get("default") is True

    def test_dev_mode(self):
        assert self._yaml()["targets"]["dev"]["mode"] == "development"

    def test_prod_mode(self):
        assert self._yaml()["targets"]["prod"]["mode"] == "production"

    def test_each_target_has_catalog(self):
        for tname, tcfg in self._yaml()["targets"].items():
            assert "catalog" in tcfg.get("variables", {}), (
                f"Target '{tname}' missing catalog variable"
            )


class TestIdempotency:
    """Running the generator twice produces identical output."""

    def test_yaml_idempotent(self):
        from backends.databricks import emit

        from conf import cfg

        first = emit(cfg._data)
        second = emit(cfg._data)
        assert first["databricks.yml"] == second["databricks.yml"]

    def test_resources_idempotent(self):
        """Python resource loader produces identical structure twice."""
        from bundle.resources import load_resources

        first = load_resources()
        second = load_resources()
        # Compare job task keys as a proxy for structural equality
        t1 = [t.task_key for t in first.jobs["stack_pipelines"].tasks]
        t2 = [t.task_key for t in second.jobs["stack_pipelines"].tasks]
        assert t1 == t2
        assert set(first.schemas.keys()) == set(second.schemas.keys())
        assert set(first.volumes.keys()) == set(second.volumes.keys())