atlas-iac/testing/tests/test_hermes_execution_pool_manifests.py

"""Shared manifest and scheduling helpers for the execution-pool suites."""

from __future__ import annotations

from pathlib import Path

import yaml


ROOT = Path(__file__).parents[2]
HERMES = ROOT / "services/hermes"
FLUX = ROOT / "clusters/atlas/flux-system/applications"


def documents(path):
    """Every non-empty YAML document in one manifest file."""
    return [item for item in yaml.safe_load_all(path.read_text()) if item]


def worker():
    """The rendered-source hermes-execution-worker StatefulSet."""
    return documents(HERMES / "execution-worker-statefulset.yaml")[0]


def mediators():
    """The three per-ordinal mediator Deployments, in ordinal order."""
    return [
        item
        for item in documents(HERMES / "execution-mediator.yaml")
        if item["kind"] == "Deployment"
    ]


def node(name, *, ready=True, cordoned=False):
    """One simulated pool-eligible worker node."""
    return {
        "name": name, "ready": ready, "cordoned": cordoned,
        "labels": {
            "kubernetes.io/arch": "arm64",
            "node-role.kubernetes.io/worker": "true",
            "kubernetes.io/hostname": name,
        },
    }


def matches(candidate, expressions):
    """Whether one node satisfies a nodeAffinity match-expression list."""
    for expression in expressions:
        actual = candidate["labels"].get(expression["key"])
        values = expression.get("values", [])
        if expression["operator"] == "In" and actual not in values:
            return False
        if expression["operator"] == "NotIn" and actual in values:
            return False
    return True


def placeable(pod, nodes, *, occupied=()):
    """Nodes this Pod spec could be scheduled onto, honouring hard constraints."""
    affinity = pod.get("affinity", {})
    required = affinity.get("nodeAffinity", {}).get(
        "requiredDuringSchedulingIgnoredDuringExecution", {}
    )
    terms = required.get("nodeSelectorTerms", [{"matchExpressions": []}])
    hard_pod = affinity.get("podAffinity", {}).get(
        "requiredDuringSchedulingIgnoredDuringExecution", []
    )
    anti = affinity.get("podAntiAffinity", {}).get(
        "requiredDuringSchedulingIgnoredDuringExecution", []
    )
    allowed = []
    for candidate in nodes:
        if not candidate["ready"] or candidate["cordoned"]:
            continue
        if not any(matches(candidate, term["matchExpressions"]) for term in terms):
            continue
        if anti and candidate["name"] in occupied:
            continue
        if hard_pod:
            # A hard podAffinity can only be satisfied where its target already is.
            continue
        allowed.append(candidate["name"])
    return allowed