homelab/ci/images/authentik-gate-sync/test_controller.py

import copy
import unittest

import controller as c

LLAMA = {
    "metadata": {
        "name": "llama-server", "namespace": "ai", "uid": "u-1",
        "labels": {"dns.renner.uno/public": "true", "kustomize.toolkit.fluxcd.io/name": "llama-server"},
        "annotations": {
            "cert-manager.io/cluster-issuer": "letsencrypt-cluster",
            "auth.renner.uno/login-flow": "sso-authentication",
            "auth.renner.uno/skip-path": r"^/(manifest\.webmanifest|sw\.js)$",
            "traefik.ingress.kubernetes.io/router.middlewares": "ingress-foo@kubernetescrd",
        },
    },
    "spec": {
        "ingressClassName": "authentik-gate",
        "tls": [{"hosts": ["llm.renner.uno"], "secretName": "llama-server-tls"}],
        "rules": [{"host": "llm.renner.uno", "http": {"paths": [
            {"path": "/", "pathType": "Prefix",
             "backend": {"service": {"name": "llama-server", "port": {"number": 8080}}}}]}}],
    },
}

SERVICE = {"spec": {"selector": {"app.kubernetes.io/name": "llama-server"},
                    "ports": [{"name": "http", "port": 8080, "targetPort": "http"}]}}


def ing(ns, name):
    return {"metadata": {"namespace": ns, "name": name}}


class Names(unittest.TestCase):
    def test_unambiguous(self):
        self.assertNotEqual(c.gate_name(ing("a-b", "c")), c.gate_name(ing("a", "b-c")))

    def test_stable_and_readable(self):
        self.assertEqual(c.gate_name(ing("ai", "llama-server")), c.gate_name(ing("ai", "llama-server")))
        self.assertTrue(c.gate_name(ing("ai", "llama-server")).startswith("ai-llama-server-"))

    def test_bounded(self):
        n = c.gate_name(ing("n" * 63, "i" * 253))
        self.assertLessEqual(len(c.outpost_service({"name": n})), 63)
        self.assertNotIn("--", n)

    def test_group_expression_is_a_literal(self):
        expr = c.group_expression("x'); import os; os.system('id'); ('")
        self.assertEqual(expr, "return ak_is_group_member(request.user, name="
                               + repr("x'); import os; os.system('id'); ('") + ")")
        compile(expr.replace("return ", "_ = "), "<expr>", "exec")


class GateSpec(unittest.TestCase):
    def test_llama(self):
        g = c.gate_spec(LLAMA)
        self.assertEqual(g["name"], c.gate_name(LLAMA))
        self.assertEqual((g["host"], g["scheme"], g["service"], g["port"]),
                         ("llm.renner.uno", "https", "llama-server", 8080))
        self.assertIsNone(g["group"])
        self.assertEqual(g["login_flow"], "sso-authentication")
        self.assertEqual(g["authorization_flow"], c.DEFAULT_AUTHZ_FLOW)

    def test_group_and_flow_override(self):
        ing = copy.deepcopy(LLAMA)
        ing["metadata"]["annotations"]["auth.renner.uno/group"] = "Home Assistant Config Editors"
        ing["metadata"]["annotations"]["auth.renner.uno/authorization-flow"] = "implicit"
        g = c.gate_spec(ing)
        self.assertEqual(g["group"], "Home Assistant Config Editors")
        self.assertEqual(g["authorization_flow"], "implicit")

    def test_rejects_two_hosts(self):
        ing = copy.deepcopy(LLAMA)
        ing["spec"]["rules"].append(copy.deepcopy(ing["spec"]["rules"][0]))
        ing["spec"]["rules"][1]["host"] = "other.renner.uno"
        with self.assertRaises(c.GateError):
            c.gate_spec(ing)

    def test_rejects_two_backends(self):
        ing = copy.deepcopy(LLAMA)
        paths = ing["spec"]["rules"][0]["http"]["paths"]
        paths.append({"path": "/api", "pathType": "Prefix",
                      "backend": {"service": {"name": "other", "port": {"number": 80}}}})
        with self.assertRaises(c.GateError):
            c.gate_spec(ing)

    def test_rejects_wildcard(self):
        ing = copy.deepcopy(LLAMA)
        ing["spec"]["rules"][0]["host"] = "*.renner.uno"
        with self.assertRaises(c.GateError):
            c.gate_spec(ing)


class Backend(unittest.TestCase):
    def test_named_target_port_stays_named(self):
        g = c.gate_spec(LLAMA)
        self.assertEqual(c.resolve_backend(g, SERVICE),
                         (8080, {"app.kubernetes.io/name": "llama-server"}, "http"))
        cnp = c.backend_ingress_cnp(g, {"app.kubernetes.io/name": "llama-server"}, "http")
        self.assertEqual(cnp["spec"]["ingress"][0]["toPorts"][0]["ports"][0]["port"], "http")

    def test_numeric_target_port(self):
        g = c.gate_spec(LLAMA)
        svc = copy.deepcopy(SERVICE)
        svc["spec"]["ports"][0]["targetPort"] = 9090
        self.assertEqual(c.resolve_backend(g, svc)[2], 9090)

    def test_missing_target_port_defaults_to_port(self):
        g = c.gate_spec(LLAMA)
        svc = copy.deepcopy(SERVICE)
        del svc["spec"]["ports"][0]["targetPort"]
        self.assertEqual(c.resolve_backend(g, svc)[2], 8080)

    def test_service_without_selector_fails(self):
        with self.assertRaises(c.GateError):
            c.resolve_backend(c.gate_spec(LLAMA), {"spec": {"ports": SERVICE["spec"]["ports"]}})

    def test_missing_port_fails(self):
        g = c.gate_spec(LLAMA)
        g["port"] = 1234
        with self.assertRaises(c.GateError):
            c.resolve_backend(g, SERVICE)


class Objects(unittest.TestCase):
    def setUp(self):
        self.g = c.gate_spec(LLAMA)

    def test_route_ingress(self):
        r = c.route_ingress(LLAMA, self.g)
        self.assertEqual(r["metadata"]["name"], "llama-server-gate")
        self.assertEqual(r["spec"]["ingressClassName"], "traefik")
        self.assertEqual(r["spec"]["tls"], LLAMA["spec"]["tls"])
        backend = r["spec"]["rules"][0]["http"]["paths"][0]["backend"]["service"]
        self.assertEqual(backend, {"name": "llama-server-gate", "port": {"number": 9000}})
        labels, ann = r["metadata"]["labels"], r["metadata"]["annotations"]
        self.assertEqual(labels["dns.renner.uno/public"], "true")
        self.assertNotIn("kustomize.toolkit.fluxcd.io/name", labels)
        self.assertEqual(labels["app.kubernetes.io/managed-by"], c.MANAGER)
        self.assertEqual(ann, {"traefik.ingress.kubernetes.io/router.middlewares": "ingress-foo@kubernetescrd"})
        self.assertEqual(r["metadata"]["ownerReferences"][0]["uid"], "u-1")

    def test_external_name(self):
        s = c.external_name_service(self.g)
        self.assertEqual(s["spec"]["externalName"], f"ak-outpost-{self.g['name']}.authentik.svc.cluster.local")

    def test_cnps(self):
        pods = {"app.kubernetes.io/name": "llama-server"}
        e = c.outpost_egress_cnp(self.g, pods, 8080)
        self.assertEqual(e["metadata"]["namespace"], "authentik")
        self.assertEqual(e["spec"]["endpointSelector"]["matchLabels"],
                         {"goauthentik.io/outpost-name": self.g["name"]})
        self.assertEqual(e["spec"]["egress"][0]["toEndpoints"][0]["matchLabels"],
                         {"k8s:io.kubernetes.pod.namespace": "ai", "app.kubernetes.io/name": "llama-server"})
        i = c.backend_ingress_cnp(self.g, pods, 8080)
        self.assertEqual(i["metadata"]["namespace"], "ai")
        self.assertEqual(i["spec"]["ingress"][0]["fromEndpoints"][0]["matchLabels"],
                         {"k8s:io.kubernetes.pod.namespace": "authentik",
                          "goauthentik.io/outpost-name": self.g["name"]})
        self.assertEqual(i["spec"]["ingress"][0]["toPorts"][0]["ports"][0]["port"], "8080")

    def test_provider_body(self):
        flows = {c.DEFAULT_AUTHZ_FLOW: "az", "sso-authentication": "au",
                 c.INVALIDATION_FLOW: "inv"}
        p = c.provider_body(self.g, 8080, flows)
        self.assertEqual(p["external_host"], "https://llm.renner.uno")
        self.assertEqual(p["internal_host"], "http://llama-server.ai.svc.cluster.local:8080")
        self.assertEqual((p["authorization_flow"], p["authentication_flow"], p["invalidation_flow"]),
                         ("az", "au", "inv"))
        self.assertEqual(p["skip_path_regex"], r"^/(manifest\.webmanifest|sw\.js)$")

    def test_status_mirrors_route_address(self):
        addrs = [{"ip": "10.2.99.86"}, {"ip": "fd99::86"}]
        route = {"status": {"loadBalancer": {"ingress": addrs}}}
        self.assertEqual(c.status_patch(LLAMA, c.lb_ingress(route)),
                         {"status": {"loadBalancer": {"ingress": addrs}}})
        synced = dict(LLAMA, status={"loadBalancer": {"ingress": addrs}})
        self.assertIsNone(c.status_patch(synced, addrs))

    def test_status_cleared_without_route(self):
        stale = dict(LLAMA, status={"loadBalancer": {"ingress": [{"ip": "10.2.99.86"}]}})
        self.assertEqual(c.status_patch(stale, []), {"status": {"loadBalancer": {"ingress": None}}})
        self.assertIsNone(c.status_patch(LLAMA, []))

    def test_drift(self):
        self.assertEqual(c.drift({"a": 1, "b": 2}, {"a": 1, "b": 3}), {"b": 3})


class FakeAuthentik:
    """In-memory stand-in for the endpoints the controller uses; records calls."""

    def __init__(self):
        self.db = {"providers/proxy/": {}, "core/applications/": {}, "outposts/instances/": {},
                   "policies/bindings/": {}, "policies/expression/": {}, "core/groups/": {}}
        self.calls = []
        self.next = 1

    def add(self, path, obj):
        obj.setdefault("pk", self.next)
        self.next += 1
        self.db[path][obj["pk"]] = obj
        return obj

    def _split(self, path):
        for coll in self.db:
            if path.startswith(coll):
                rest = path[len(coll):].strip("/")
                return coll, rest or None
        raise AssertionError(path)

    def _find(self, coll, key):
        if coll == "core/applications/":
            return next((a for a in self.db[coll].values() if a["slug"] == key), None)
        return self.db[coll].get(int(key) if str(key).isdigit() else key)

    def list(self, path, **q):
        out = []
        for o in self.db[path].values():
            ok = True
            for k, v in q.items():
                field = k.replace("__iexact", "")
                ok &= str(o.get(field)).lower() == str(v).lower() if "__iexact" in k else o.get(field) == v
            if ok:
                out.append(copy.deepcopy(o))
        return out

    def one(self, path, **q):
        found = self.list(path, **q)
        return found[0] if found else None

    def get(self, path):
        coll, key = self._split(path)
        o = self._find(coll, key)
        return copy.deepcopy(o) if o else None

    def req(self, method, path, body=None):
        self.calls.append((method, path))
        coll, key = self._split(path)
        if method == "POST":
            obj = self.add(coll, dict(body))
            if coll == "core/applications/":
                obj["pbm_uuid"] = f"pbm-{obj['pk']}"
                self.db["providers/proxy/"][obj["provider"]]["assigned_application_slug"] = obj["slug"]
            return copy.deepcopy(obj)
        obj = self._find(coll, key)
        if method == "PATCH":
            # ProxyProviderSerializer.validate sees only the request body.
            if coll == "providers/proxy/" and body.get("mode", "proxy") == "proxy" \
                    and body.get("internal_host", "") == "":
                raise AssertionError("400 internal_host")
            obj.update(body)
            return copy.deepcopy(obj)
        if method == "DELETE":
            del self.db[coll][obj["pk"]]
            if coll == "core/applications/":
                self.db["providers/proxy/"][obj["provider"]]["assigned_application_slug"] = None
                for b in [b for b in self.db["policies/bindings/"].values() if b["target"] == obj["pbm_uuid"]]:
                    del self.db["policies/bindings/"][b["pk"]]
            return None
        raise AssertionError(method)


FLOWS = {c.DEFAULT_AUTHZ_FLOW: "az", "sso-authentication": "au", c.INVALIDATION_FLOW: "inv"}


class Reconcile(unittest.TestCase):
    def setUp(self):
        self.ak = FakeAuthentik()
        self.g = c.gate_spec(LLAMA)

    def gate(self, group=None):
        g = dict(self.g, group=group)
        c.ensure_gate(self.ak, g, 8080, FLOWS, "conn")
        return g

    def test_creates_owned_gate(self):
        self.gate()
        app = self.ak.get(f"core/applications/{self.g['name']}/")
        self.assertEqual(app["meta_publisher"], c.MANAGER)
        outpost = self.ak.one("outposts/instances/", name__iexact=self.g["name"])
        self.assertEqual(outpost["providers"], [app["provider"]])
        self.assertEqual(self.ak.list("policies/bindings/"), [])

    def test_second_pass_is_a_noop(self):
        self.gate()
        before = len(self.ak.calls)
        self.gate()
        self.assertEqual(self.ak.calls[before:], [])

    def test_flow_change_updates_provider(self):
        self.gate()
        self.g = dict(self.g, login_flow=None)
        self.gate()
        provider = self.ak.one("providers/proxy/", name__iexact=self.g["name"])
        self.assertIsNone(provider["authentication_flow"])

    def test_ignores_foreign_provider_names(self):
        self.ak.add("providers/proxy/", {"name": "aaa-first-alphabetically", "assigned_application_slug": "x"})
        self.gate()
        self.assertEqual(len(self.ak.list("providers/proxy/")), 2)

    def test_refuses_unowned_application(self):
        self.ak.add("core/applications/", {"slug": self.g["name"], "meta_publisher": "", "provider": None})
        with self.assertRaises(c.GateError):
            self.gate()

    def test_missing_group_is_a_gate_error(self):
        with self.assertRaises(c.GateError):
            self.gate(group="Nope")

    def test_group_change_binds_before_unbinding(self):
        self.ak.add("core/groups/", {"name": "A"})
        self.ak.add("core/groups/", {"name": "B"})
        self.gate(group="A")
        bindings = self.ak.list("policies/bindings/")
        self.assertEqual(len(bindings), 1)
        policy = self.ak.one("policies/expression/", name=self.g["name"])
        self.assertIn("'A'", policy["expression"])
        self.gate(group="B")
        policy = self.ak.one("policies/expression/", name=self.g["name"])
        self.assertIn("'B'", policy["expression"])
        self.assertEqual([b["pk"] for b in self.ak.list("policies/bindings/")], [bindings[0]["pk"]])

    def test_dropping_group_removes_binding_after_nothing_new(self):
        self.ak.add("core/groups/", {"name": "A"})
        self.gate(group="A")
        self.gate(group=None)
        self.assertEqual(self.ak.list("policies/bindings/"), [])
        self.assertIsNone(self.ak.one("policies/expression/", name=self.g["name"]))

    def test_remove_gate_and_ownership_scan(self):
        self.ak.add("providers/proxy/", {"name": "foreign", "assigned_application_slug": "foreign"})
        self.ak.add("core/applications/", {"slug": "foreign", "meta_publisher": "", "provider": 1})
        self.gate()
        owned = list(c.owned_applications(self.ak))
        self.assertEqual([a["slug"] for a in owned], [self.g["name"]])
        c.remove_gate(self.ak, owned[0])
        self.assertIsNone(self.ak.get(f"core/applications/{self.g['name']}/"))
        self.assertEqual([p["name"] for p in self.ak.list("providers/proxy/")], ["foreign"])
        self.assertEqual(self.ak.list("outposts/instances/"), [])

    def test_adopts_provider_orphaned_by_interrupted_teardown(self):
        self.gate()
        app = self.ak.get(f"core/applications/{self.g['name']}/")
        self.ak.req("DELETE", f"core/applications/{self.g['name']}/")
        self.gate()
        self.assertEqual(self.ak.get(f"core/applications/{self.g['name']}/")["provider"], app["provider"])


if __name__ == "__main__":
    unittest.main()