import types
import unittest
from unittest.mock import Mock, patch

from sglang.srt.disaggregation.common import conn as common_conn


class TestDisaggregationRuntimeFingerprint(unittest.TestCase):
    def _manager(self):
        mgr = common_conn.CommonKVManager.__new__(common_conn.CommonKVManager)
        mgr.prefill_info_table = {}
        mgr.kv_args = types.SimpleNamespace(page_size=64, engine_rank=0)
        mgr.server_args = types.SimpleNamespace(kv_cache_dtype="fp8_e4m3")
        mgr._resolve_rank_mapping = Mock()
        return mgr

    def _response(self, payload, status_code=200):
        response = Mock()
        response.status_code = status_code
        response.json.return_value = payload
        response.text = ""
        return response

    def _payload(self, **overrides):
        payload = {
            "attn_tp_size": 8,
            "attn_cp_size": 8,
            "dp_size": 1,
            "pp_size": 1,
            "page_size": 64,
            "kv_cache_dtype": "fp8_e4m3",
            "follow_bootstrap_room": True,
        }
        payload.update(overrides)
        return payload

    def test_try_ensure_parallel_info_rejects_runtime_source_fingerprint_mismatch(self):
        mgr = self._manager()
        with patch.object(
            common_conn.requests,
            "get",
            return_value=self._response(
                self._payload(
                    runtime_source_fingerprint="prefill-source-hash",
                    runtime_source_root="/sgl-workspace/sglang-tai",
                )
            ),
        ), patch.object(
            common_conn,
            "get_runtime_source_fingerprint",
            return_value="decode-source-hash",
        ), patch.object(
            common_conn,
            "get_runtime_source_root",
            return_value="/sgl-workspace/sglang",
        ):
            with self.assertRaisesRegex(
                RuntimeError,
                "Runtime source fingerprint mismatch",
            ):
                mgr.try_ensure_parallel_info("g0034:18992")

        mgr._resolve_rank_mapping.assert_not_called()
        self.assertEqual(mgr.prefill_info_table, {})

    def test_try_ensure_parallel_info_accepts_matching_runtime_source_fingerprint(self):
        mgr = self._manager()
        with patch.object(
            common_conn.requests,
            "get",
            return_value=self._response(
                self._payload(
                    runtime_source_fingerprint="same-source-hash",
                    runtime_source_root="/sgl-workspace/sglang-tai",
                )
            ),
        ), patch.object(
            common_conn,
            "get_runtime_source_fingerprint",
            return_value="same-source-hash",
        ), patch.object(
            common_conn,
            "get_runtime_source_root",
            return_value="/sgl-workspace/sglang-tai",
        ):
            self.assertTrue(mgr.try_ensure_parallel_info("g0034:18992"))

        mgr._resolve_rank_mapping.assert_called_once()
        self.assertIn("g0034:18992", mgr.prefill_info_table)


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