example_utils.py

"""Utilities for working with vendored example manifests in tests."""

from __future__ import annotations

import os
import subprocess
from functools import lru_cache
from pathlib import Path

import yaml

from helpers import kubectl, wait_for, wait_for_stateful_set_ready


def _discover_repo_root(start: Path) -> Path:
    for candidate in (start, *start.parents):
        if (candidate / "examples" / "recommendations.json").is_file():
            return candidate
    raise RuntimeError("unable to locate repo root from example_utils.py")


_EXAMPLES_ROOT_OVERRIDE = os.environ.get("EXAMPLES_ROOT")
if _EXAMPLES_ROOT_OVERRIDE:
    EXAMPLES_ROOT = Path(_EXAMPLES_ROOT_OVERRIDE).resolve()
    REPO_ROOT = EXAMPLES_ROOT.parent
else:
    REPO_ROOT = _discover_repo_root(Path(__file__).resolve().parent)
    EXAMPLES_ROOT = (REPO_ROOT / "examples").resolve()
INVALID_EXAMPLES_ROOT = EXAMPLES_ROOT / "invalid"
EXCLUDED_VALID_EXAMPLE_ROOTS = {
    EXAMPLES_ROOT / "gpus",
}

OPTIONAL_API_GROUPS = {
    "autoscaling.k8s.io": [
        EXAMPLES_ROOT / "automationstrategy" / "vpa-filter.yaml",
        EXAMPLES_ROOT / "automationstrategy" / "vpa-filter-default.yaml",
        EXAMPLES_ROOT / "automationstrategy" / "vpa-filter-explicit-containers.yaml",
        EXAMPLES_ROOT / "automationstrategy" / "vpa-filter-recommendation-managed.yaml",
        EXAMPLES_ROOT / "automationstrategy" / "vpa-filter-prevpa.yaml",
    ],
    "keda.sh": [
        EXAMPLES_ROOT / "staticpolicy" / "with-keda-hpa-filter.yaml",
    ],
}


def all_valid_example_manifests() -> list[Path]:
    return sorted(
        path
        for path in EXAMPLES_ROOT.rglob("*.yaml")
        if INVALID_EXAMPLES_ROOT not in path.parents
        and not any(excluded in path.parents for excluded in EXCLUDED_VALID_EXAMPLE_ROOTS)
    )


def all_invalid_example_manifests() -> list[Path]:
    return sorted(INVALID_EXAMPLES_ROOT.rglob("*.yaml"))


def manifest_documents(manifest_path: Path) -> list[dict]:
    with manifest_path.open() as handle:
        return [
            doc
            for doc in yaml.safe_load_all(handle)
            if doc and doc.get("kind") and doc.get("metadata", {}).get("name")
        ]


@lru_cache(maxsize=None)
def api_group_available(kube_context: str, api_group: str) -> bool:
    output = kubectl(
        "api-resources",
        "--api-group",
        api_group,
        "-o",
        "name",
        context=kube_context,
    )
    return bool(output.strip())


def skip_reason(path: Path, kube_context: str) -> str | None:
    for api_group, manifests in OPTIONAL_API_GROUPS.items():
        if path in manifests and not api_group_available(kube_context, api_group):
            return f"optional API group {api_group} is not installed in the cluster"

    if path == EXAMPLES_ROOT / "automationstrategy" / "hpa-filter-container.yaml":
        result = subprocess.run(
            [
                "kubectl",
                "--context",
                kube_context,
                "explain",
                "--api-version=autoscaling/v2",
                "horizontalpodautoscaler.spec.metrics.containerResource",
            ],
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            detail = result.stderr.strip()
            if detail:
                return (
                    "cluster does not support the HPA containerResource metric used by this example: "
                    f"{detail}"
                )
            return "cluster does not support the HPA containerResource metric used by this example"
    return None


def apply_manifest(manifest_path: Path, kube_context: str) -> None:
    kubectl("apply", "-f", str(manifest_path), context=kube_context)


def delete_manifest_in_reverse(manifest_path: Path, kube_context: str) -> None:
    """Delete manifest resources in reverse document order."""
    for doc in reversed(manifest_documents(manifest_path)):
        kind = doc["kind"]
        name = doc["metadata"]["name"]
        namespace = doc["metadata"].get("namespace")
        cmd = [
            "kubectl",
            "--context",
            kube_context,
            "delete",
            kind,
            name,
            "--ignore-not-found",
            "--wait=false",
        ]
        if namespace:
            cmd += ["-n", namespace]
        subprocess.run(cmd, capture_output=True)
        if kind == "Namespace":
            wait_for(
                lambda: not subprocess.run(
                    [
                        "kubectl",
                        "--context",
                        kube_context,
                        "get",
                        "namespace",
                        name,
                    ],
                    capture_output=True,
                    text=True,
                ).returncode
                == 0,
                timeout=60,
                message=f"namespace {name} deletion",
            )


def assert_declared_resources_exist(manifest_path: Path, kube_context: str) -> None:
    for doc in manifest_documents(manifest_path):
        kind = doc["kind"]
        name = doc["metadata"]["name"]
        namespace = doc["metadata"].get("namespace")
        args = ["get", kind, name, "-o", "name"]
        if namespace:
            args += ["-n", namespace]
        kubectl(*args, context=kube_context)


def wait_for_declared_workloads_ready(manifest_path: Path, k8s_clients) -> None:
    def wait_for_ready(read_workload, kind: str, namespace: str, name: str, expected_replicas: int):
        wait_for(
            lambda: (
                (workload := read_workload(name, namespace))
                and (workload.status.ready_replicas or 0) >= expected_replicas
            ),
            timeout=180,
            message=f"{kind.lower()} {namespace}/{name} readiness",
        )

    for doc in manifest_documents(manifest_path):
        kind = doc["kind"]
        namespace = doc["metadata"].get("namespace")
        name = doc["metadata"]["name"]
        expected_replicas = doc.get("spec", {}).get("replicas", 1)

        if kind == "Deployment":
            wait_for_ready(
                k8s_clients.apps.read_namespaced_deployment,
                kind,
                namespace,
                name,
                expected_replicas,
            )
            continue

        if kind == "StatefulSet":
            wait_for_stateful_set_ready(k8s_clients.apps, namespace, name, expected_replicas)