import unittest
from agents.mcp_server.server import _sanitize_opax_deployment_status

class TestSanitizeOpaxDeploymentStatus(unittest.TestCase):

    def test_valid_inputs(self):
        """Tests valid integer/string generations and a valid UTC timestamp."""
        service_data = {
            "generation": 10,
            "observedGeneration": "9",
            "updateTime": "2023-10-27T10:00:00Z"
        }
        result = _sanitize_opax_deployment_status(service_data)
        self.assertEqual(result["generation"], 10)
        self.assertEqual(result["observed_generation"], 9)
        self.assertEqual(result["last_update_time"], "2023-10-27T10:00:00Z")
        self.assertIsNone(result["reason_code"])

    def test_zero_values(self):
        """Tests handling of zero for generation values."""
        service_data = {"generation": 0, "observedGeneration": "0"}
        result = _sanitize_opax_deployment_status(service_data)
        self.assertEqual(result["generation"], 0)
        self.assertEqual(result["observed_generation"], 0)

    def test_invalid_generation_values(self):
        """Tests various invalid generation inputs."""
        invalid_inputs = [
            True, False, -1, 1.5, "-1", " 1 ", "abc", "١", None, {}, []
        ]
        for val in invalid_inputs:
            with self.subTest(val=val):
                res1 = _sanitize_opax_deployment_status({"generation": val})
                res2 = _sanitize_opax_deployment_status({"observedGeneration": val})
                self.assertIsNone(res1["generation"])
                self.assertIsNone(res2["observed_generation"])

    def test_fractional_second_timestamp_is_truncated(self):
        """Tests that fractional seconds are correctly truncated."""
        service_data = {"updateTime": "2023-10-27T10:00:00.123456Z"}
        result = _sanitize_opax_deployment_status(service_data)
        self.assertEqual(result["last_update_time"], "2023-10-27T10:00:00Z")

    def test_invalid_and_non_z_timestamps(self):
        """Tests invalid, tz-naive, missing, and non-Z timestamps."""
        invalid_timestamps = [
            "2023-10-27T10:00:00",          # Missing Z (naive)
            "2023-10-27T10:00:00ZZ",         # Multiple Z
            "2023-10-27T11:00:00+01:00",  # Valid but not Z-suffix
            "2023-10-27 10:00:00Z",         # Malformed
            "not-a-timestamp",              # Invalid
            1698399600,                     # Not a string
            None,                           # Missing
        ]
        for ts in invalid_timestamps:
            with self.subTest(ts=ts):
                result = _sanitize_opax_deployment_status({"updateTime": ts})
                self.assertIsNone(result["last_update_time"])

    def test_exact_key_set_and_fixed_values(self):
        """Tests for exact key set and fixed service/region/reason_code."""
        result = _sanitize_opax_deployment_status({})
        expected_keys = {
            "service", "region", "generation", "observed_generation",
            "last_update_time", "reason_code"
        }
        self.assertEqual(set(result.keys()), expected_keys)
        self.assertEqual(result["service"], "opax-mcp")
        self.assertEqual(result["region"], "us-central1")
        self.assertIsNone(result["reason_code"])

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