import json
import unittest
from pathlib import Path

import yaml
from ansible.plugins.filter.core import combine
from jinja2 import Environment, StrictUndefined


class ReadinessTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        tasks = yaml.safe_load(
            (Path(__file__).parents[1] / "roles/k3s/tasks/main.yml").read_text()
        )
        env = Environment(undefined=StrictUndefined)
        env.filters["from_json"] = json.loads
        env.filters["combine"] = combine
        cls.env = env
        cls.tasks = tasks
        cls.conditions = {
            task["register"]: env.compile_expression(task["until"])
            for task in tasks
            if "until" in task
        }

    def test_node_must_be_ready_at_the_requested_version(self):
        condition = self.conditions["k3s_node_status"]
        for version, ready, expected in [
            ("v1.35.9+k3s1", "True", True),
            ("v1.35.6+k3s1", "True", False),
            ("v1.35.9+k3s1", "False", False),
        ]:
            with self.subTest(version=version, ready=ready):
                status = {"status": {
                    "nodeInfo": {"kubeletVersion": version},
                    "conditions": [{"type": "Ready", "status": ready}],
                }}
                self.assertEqual(condition(
                    k3s_node_status={"rc": 0, "stdout": json.dumps(status)},
                    k3s_version="v1.35.9+k3s1",
                ), expected)
        self.assertFalse(condition(
            k3s_node_status={"rc": 1, "stdout": ""},
            k3s_version="v1.35.9+k3s1",
        ))

    def test_api_requires_a_successful_readyz_response(self):
        condition = self.conditions["k3s_api_ready"]
        for rc, stdout, expected in [(0, "ok", True), (1, "ok", False), (0, "", False)]:
            with self.subTest(rc=rc, stdout=stdout):
                self.assertEqual(condition(k3s_api_ready={"rc": rc, "stdout": stdout}), expected)

    def test_mirror_changes_restart_existing_nodes_without_a_version_upgrade(self):
        task = next(t for t in self.tasks if t["name"] == "Restart K3s after updating registry mirrors")
        condition = self.env.compile_expression(" and ".join(f"({c})" for c in task["when"]))
        for changed, exists, current, expected in [
            (True, True, "v1.35.9+k3s1", True),
            (False, True, "v1.35.9+k3s1", False),
            (True, True, "v1.35.6+k3s1", False),
            (True, False, "", False),
        ]:
            with self.subTest(changed=changed, exists=exists, current=current):
                self.assertEqual(condition(
                    k3s_registries_config={"changed": changed},
                    k3s_binary={"stat": {"exists": exists}},
                    k3s_current_version={"stdout": current} if exists else {},
                    k3s_version="v1.35.9+k3s1",
                ), expected)

    def test_host_mirror_override_preserves_other_registries(self):
        template = self.env.from_string(
            (Path(__file__).parents[1] / "roles/k3s/templates/registries.yaml.j2").read_text()
        )
        mirrors = {"docker.io": ["docker.1ms.run", "docker.m.daocloud.io"], "ghcr.io": ["ghcr.m.daocloud.io"]}
        for override in [{}, {"docker.io": list(reversed(mirrors["docker.io"]))}]:
            result = yaml.safe_load(template.render(
                registry_mirrors=mirrors, registry_mirror_overrides=override,
            ))["mirrors"]
            self.assertEqual(result["ghcr.io"]["endpoint"], ["https://ghcr.m.daocloud.io"])
            self.assertEqual(result["docker.io"]["endpoint"], [
                "https://" + host for host in override.get("docker.io", mirrors["docker.io"])
            ])


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