import json
import unittest
from pathlib import Path

import yaml
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
        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)


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