from __future__ import annotations

import unittest
from unittest.mock import call, patch

from monitors.yumi import service_ops


class NacosDrainStateTests(unittest.TestCase):
    def test_enabled_state_uses_catalog_when_disabled_instance_is_filtered_from_instance_list(self) -> None:
        service = {
            "host": "app-1",
            "service": "other",
            "ip": "10.2.11.3",
            "port": 2400,
            "kind": "java",
        }

        with (
            patch.object(service_ops, "current_namespace", return_value=("namespace-id", "namespace")),
            patch.object(
                service_ops,
                "list_instances",
                return_value=[
                    {
                        "ip": "10.2.12.6",
                        "port": 2400,
                        "healthy": True,
                        "enabled": True,
                    }
                ],
            ),
            patch.object(
                service_ops,
                "catalog_instances",
                return_value=[
                    {
                        "ip": "10.2.11.3",
                        "port": 2400,
                        "healthy": True,
                        "enabled": False,
                    }
                ],
            ),
        ):
            result = service_ops.wait_java_registration_enabled_state(service, False, 1)

        self.assertTrue(result["checked"])
        self.assertFalse(result["enabled"])
        self.assertTrue(result["healthy"])
        self.assertEqual(result["source"], "catalog/instances")

    def test_drain_restores_enabled_when_wait_state_fails_after_update(self) -> None:
        service = {
            "host": "app-1",
            "service": "other",
            "ip": "10.2.11.3",
            "port": 2400,
            "kind": "java",
        }

        with (
            patch.object(service_ops, "should_drain_java_registration", return_value=True),
            patch.object(service_ops, "wait_java_registration_enabled_state", side_effect=RuntimeError("not ready")),
            patch.object(service_ops, "update_java_registration_enabled", return_value={"checked": True}) as update_mock,
        ):
            with self.assertRaises(RuntimeError):
                service_ops.drain_java_registration(service)

        self.assertEqual(
            update_mock.call_args_list,
            [
                call(service, False),
                call(service, True),
            ],
        )


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