Model/tests/test_lambda_packaging.py

"""Static packaging linter for Lambda container images.

Every Lambda here ships as a Docker image that copies a *subset* of the repo
(``COPY utils/ utils/``, ``COPY backend/ backend/``, ...) and then runs a
handler via ``CMD ["<module>.handler"]``. If the handler's import graph reaches
a top-level package the Dockerfile forgot to ``COPY``, the function dies at cold
start with ``Runtime.ImportModuleError: No module named '<pkg>'`` — but only
*inside the image*. In the dev/test tree every package is present, so a plain
``import`` test can't see the gap. This is exactly how ``No module named
'domain'`` reached a deployed address2UPRN.

The RIE smoke tests (.github/workflows/_smoke_test_lambda.yml) catch this too,
but only by building the full image (minutes) and only for hand-listed services.
This test catches the same class of bug in milliseconds, locally, for *every*
handler Dockerfile — by statically computing each handler's import-time module
graph and asserting every repo file it reaches is copied into the image.

Scope: import-time (module-level) imports only — the ones that run at Lambda
init, which is what ImportModuleError is about. Imports inside function bodies
and under ``if TYPE_CHECKING:`` are deliberately ignored. Third-party / stdlib
imports are out of scope (that's requirements.txt's job, covered by the RIE
smoke test actually installing and importing).
"""

from __future__ import annotations

import ast
import json
import re
from pathlib import Path
from typing import Optional

import pytest

REPO_ROOT = Path(__file__).resolve().parents[1]

# Dockerfiles that are not Lambda handlers (test harness, dev containers).
_SKIP_DOCKERFILES = {"Dockerfile.test", "Dockerfile.test.dockerignore"}
_SKIP_PARTS = {".git", "node_modules", ".devcontainer"}


def _toplevel_names() -> set[str]:
    """Top-level repo packages/modules — the namespace handler imports resolve
    against (imports are absolute: ``domain.x``, ``backend.y``)."""
    names: set[str] = set()
    for p in REPO_ROOT.iterdir():
        if p.name.startswith(".") or p.name == "__pycache__":
            continue
        if p.is_dir():
            names.add(p.name)
        elif p.suffix == ".py":
            names.add(p.stem)
    return names


_TOP = _toplevel_names()


def _is_type_checking(test: ast.expr) -> bool:
    if isinstance(test, ast.Name):
        return test.id == "TYPE_CHECKING"
    if isinstance(test, ast.Attribute):
        return test.attr == "TYPE_CHECKING"
    return False


def _file_package_parts(path: Path) -> list[str]:
    """The components of ``__package__`` Python assigns when importing ``path``.

    For a regular module ``a/b/c.py`` and for a package ``a/b/__init__.py`` alike
    this is the containing directory (``["a", "b"]``) — i.e. the anchor that
    ``from . import x`` resolves against."""
    return list(path.relative_to(REPO_ROOT).parts)[:-1]


def _import_time_imports(path: Path) -> list[str]:
    """Absolute module names imported when ``path`` is imported (i.e. at Lambda
    init). Descends into module-level if/try/with and class bodies, but not into
    function bodies (lazy) or ``if TYPE_CHECKING:`` blocks (never executed).

    Relative imports (``from .x import y``) are resolved to their absolute name
    against ``path``'s package — the codebase re-exports through package
    ``__init__.py`` files this way, so dropping them would hide real init-time
    dependencies (e.g. ``functions/__init__.py`` -> ``from .portfolio_functions
    import *`` -> ... -> ``infrastructure``)."""
    try:
        tree = ast.parse(path.read_text(encoding="utf-8"), str(path))
    except (SyntaxError, UnicodeDecodeError):
        return []
    pkg_parts = _file_package_parts(path)
    out: list[str] = []

    def _relative_base(level: int) -> list[str]:
        # level 1 anchors on the package itself; each extra level climbs one up.
        keep = len(pkg_parts) - (level - 1)
        return pkg_parts[:keep] if keep > 0 else []

    def visit(stmts: list[ast.stmt]) -> None:
        for node in stmts:
            if isinstance(node, ast.Import):
                out.extend(alias.name for alias in node.names)
            elif isinstance(node, ast.ImportFrom):
                if not node.level:  # absolute import
                    if node.module:
                        out.append(node.module)
                else:  # relative import — resolve against this file's package
                    base = _relative_base(node.level)
                    if node.module:  # from .pkg.mod import name
                        out.append(".".join(base + node.module.split(".")))
                    else:  # from . import a, b  -> base.a, base.b (submodules)
                        out.extend(".".join(base + [alias.name]) for alias in node.names)
            elif isinstance(node, ast.If):
                if _is_type_checking(node.test):
                    continue
                visit(node.body)
                visit(node.orelse)
            elif isinstance(node, ast.Try):
                visit(node.body)
                visit(node.orelse)
                visit(node.finalbody)
                for handler in node.handlers:
                    visit(handler.body)
            elif isinstance(node, ast.With):
                visit(node.body)
            elif isinstance(node, ast.ClassDef):
                visit(node.body)
            # FunctionDef / AsyncFunctionDef bodies are intentionally skipped.

    visit(tree.body)
    return out


def _module_files(module: str) -> list[Path]:
    """Every repo file executed when ``module`` is imported: the module's own
    file *plus* each ancestor package's ``__init__.py``.

    Importing ``a.b.c`` runs ``a/__init__.py``, ``a/b/__init__.py`` and
    ``a/b/c.py`` (or ``a/b/c/__init__.py``) in turn — so an ``__init__.py`` part
    way down the path can pull in a whole subtree (and the package it lives in
    must be COPYed). ``_module_to_file`` resolves only the leaf, which is why the
    closure used to stop short of those intermediate packages."""
    parts = module.split(".")
    files: list[Path] = []
    for depth in range(1, len(parts) + 1):
        base = REPO_ROOT.joinpath(*parts[:depth])
        init = base / "__init__.py"
        if init.is_file():
            files.append(init)
        if depth == len(parts):  # the leaf may be a plain module file
            leaf = base.with_suffix(".py")
            if leaf.is_file():
                files.append(leaf)
    return files


def _import_closure(start: Path) -> dict[Path, Optional[Path]]:
    """Repo files reachable from ``start`` via import-time imports, mapped to the
    first file that imported each (for blame in failure messages)."""
    reached: dict[Path, Optional[Path]] = {}
    stack: list[tuple[Path, Optional[Path]]] = [(start, None)]
    while stack:
        path, importer = stack.pop()
        if path in reached:
            continue
        reached[path] = importer
        for module in _import_time_imports(path):
            if module.split(".")[0] not in _TOP:
                continue  # stdlib / third-party — not our concern here
            for target in _module_files(module):
                if target not in reached:
                    stack.append((target, path))
    return reached


def _norm(path_token: str) -> str:
    return path_token.lstrip("./").rstrip("/")


def _parse_handler_spec(dockerfile_text: str) -> Optional[str]:
    """The ``<module>.handler`` string from the ``CMD`` line, or None if this
    isn't a Lambda handler image."""
    match = re.search(r"^CMD\s+(\[.*\]|.+)$", dockerfile_text, re.MULTILINE)
    if not match:
        return None
    raw = match.group(1).strip()
    try:
        parsed = json.loads(raw)
        spec = parsed[0] if isinstance(parsed, list) and parsed else raw
    except json.JSONDecodeError:
        spec = raw.strip('"')
    return spec if isinstance(spec, str) and spec.endswith(".handler") else None


def _parse_copies(dockerfile_text: str) -> list[tuple[list[str], str]]:
    """``COPY`` instructions as (sources, dest), dropping ``--flag`` tokens."""
    copies: list[tuple[list[str], str]] = []
    for match in re.finditer(r"^COPY\s+(.+)$", dockerfile_text, re.MULTILINE):
        tokens = [t for t in match.group(1).split() if not t.startswith("--")]
        if len(tokens) < 2:
            continue
        *sources, dest = tokens
        copies.append((sources, dest))
    return copies


def _resolve_handler_file(
    spec: str, copies: list[tuple[list[str], str]]
) -> Optional[Path]:
    """Map a handler spec to its repo source file.

    Handles both in-place layouts (``backend.foo.handler`` -> ``backend/foo/
    handler.py``, present via ``COPY backend/ backend/``) and root-placed
    handlers (``main.handler`` where a ``COPY <src> /var/task/main.py`` or
    ``COPY <src>/main.py .`` puts the file at the image root)."""
    module_path, _func = spec.rsplit(".", 1)

    direct = REPO_ROOT / (module_path.replace(".", "/") + ".py")
    if direct.is_file():
        return direct

    # Root-placed module: find the COPY whose destination basename matches.
    wanted = module_path.split("/")[-1] + ".py"
    for sources, dest in copies:
        dest_norm = _norm(dest)
        dest_is_named_file = Path(dest_norm).name == wanted
        dest_is_dir = dest_norm in ("", module_path.split("/")[-1]) or dest.endswith("/")
        for src in sources:
            src_path = REPO_ROOT / _norm(src)
            if not src_path.is_file():
                continue
            if dest_is_named_file or (dest_is_dir and src_path.name == wanted):
                return src_path
    return None


def _is_copied(rel_path: str, copies: list[tuple[list[str], str]]) -> bool:
    """Whether a repo-relative file path lands in the image via some COPY."""
    rel_path = _norm(rel_path)
    for sources, _dest in copies:
        for src in sources:
            src_norm = _norm(src)
            if src_norm == "" or src_norm == rel_path or rel_path.startswith(src_norm + "/"):
                return True
    return False


def _package_dir_present(pkg_rel: str, copies: list[tuple[list[str], str]]) -> bool:
    """Whether the image will contain ``pkg_rel`` as a directory because some
    COPY brings in a file beneath it. Used to excuse an un-copied package
    ``__init__.py``: in Python 3 a directory present without its ``__init__.py``
    imports fine as a *namespace package*, so the missing ``__init__`` is not a
    cold-start ``ModuleNotFoundError`` (only a wholly-absent package is)."""
    pkg_rel = _norm(pkg_rel)
    for sources, _dest in copies:
        for src in sources:
            src_norm = _norm(src)
            if src_norm == pkg_rel or src_norm.startswith(pkg_rel + "/"):
                return True
    return False


def _discover_handler_dockerfiles() -> list[Path]:
    found: list[Path] = []
    for path in REPO_ROOT.rglob("*Dockerfile*"):
        if path.name in _SKIP_DOCKERFILES:
            continue
        if any(part in _SKIP_PARTS for part in path.relative_to(REPO_ROOT).parts):
            continue
        try:
            text = path.read_text(encoding="utf-8")
        except OSError:
            continue
        if _parse_handler_spec(text):
            found.append(path)
    return sorted(found)


_HANDLER_DOCKERFILES = _discover_handler_dockerfiles()


def test_handler_dockerfiles_discovered() -> None:
    """Guard against the discovery silently finding nothing (e.g. a refactor
    that renames Dockerfiles), which would make every check below vacuous."""
    assert _HANDLER_DOCKERFILES, "no Lambda handler Dockerfiles found under repo root"


@pytest.mark.parametrize(
    "dockerfile",
    _HANDLER_DOCKERFILES,
    ids=[str(p.relative_to(REPO_ROOT)) for p in _HANDLER_DOCKERFILES],
)
def test_lambda_image_copies_full_import_closure(dockerfile: Path) -> None:
    """Every repo file the handler imports at init must be COPYed into the image."""
    text = dockerfile.read_text(encoding="utf-8")
    spec = _parse_handler_spec(text)
    assert spec is not None  # discovery guaranteed this
    copies = _parse_copies(text)

    handler_file = _resolve_handler_file(spec, copies)
    assert handler_file is not None, (
        f"{dockerfile.relative_to(REPO_ROOT)}: could not locate the source file "
        f"for CMD handler {spec!r}. Update _resolve_handler_file if this is a new "
        f"handler layout."
    )

    missing: list[str] = []
    for reached, importer in _import_closure(handler_file).items():
        rel = str(reached.relative_to(REPO_ROOT))
        if _is_copied(rel, copies):
            continue
        # An un-copied package __init__.py is non-fatal when its directory still
        # exists in the image (some other file under it is copied): Python falls
        # back to a namespace package. We still traverse such __init__ files for
        # their imports above; we just don't demand they be copied. A wholly
        # absent package (no file under it copied) is a real ModuleNotFoundError.
        if reached.name == "__init__.py" and _package_dir_present(
            str(reached.parent.relative_to(REPO_ROOT)), copies
        ):
            continue
        blame = (
            str(importer.relative_to(REPO_ROOT)) if importer else "(handler entrypoint)"
        )
        missing.append(f"  - {rel}\n      imported by {blame}")

    assert not missing, (
        f"{dockerfile.relative_to(REPO_ROOT)} runs `{spec}` but does not COPY "
        f"{len(missing)} file(s) it imports at init. The Lambda will fail at cold "
        f"start with Runtime.ImportModuleError. Add the missing top-level "
        f"package(s) as `COPY <pkg>/ <pkg>/`:\n" + "\n".join(sorted(missing))
    )