[PD] Refactor hybrid state transfer (#24932)

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -0,0 +1,49 @@
+import unittest
+
+import numpy as np
+
+from sglang.srt.disaggregation.common.utils import (
+    pack_int_lists,
+    pack_list_of_buffers,
+    unpack_int_lists,
+    unpack_list_of_buffers,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="stage-a-test-cpu")
+
+
+class TestDisaggregationWire(unittest.TestCase):
+    def test_int_lists_roundtrip(self):
+        cases = [
+            ("Q", [[1, 2, 3], [4]]),
+            ("I", [[10, 20], [30, 40, 50]]),
+            ("i", [[-1, 2], [3, -4, 5]]),
+        ]
+        for fmt, sample in cases:
+            packed = pack_int_lists(sample, fmt)
+            self.assertEqual(unpack_int_lists(packed, fmt), sample, msg=fmt)
+
+    def test_pack_accepts_ndarray(self):
+        arrs = [
+            np.array([1, 2, 3], dtype=np.int32),
+            np.array([4, 5], dtype=np.int32),
+        ]
+        packed = pack_int_lists(arrs, "i")
+        self.assertEqual(unpack_int_lists(packed, "i"), [[1, 2, 3], [4, 5]])
+
+    def test_empty_outer_list(self):
+        self.assertEqual(pack_int_lists([], "Q"), b"")
+        self.assertEqual(unpack_int_lists(b"", "Q"), [])
+
+    def test_empty_inner_list(self):
+        packed = pack_int_lists([[]], "I")
+        self.assertEqual(unpack_int_lists(packed, "I"), [[]])
+
+    def test_list_of_buffers_roundtrip(self):
+        bufs = [b"abc", b"", b"de", b"x" * 17]
+        self.assertEqual(unpack_list_of_buffers(pack_list_of_buffers(bufs)), bufs)
+
+
+if __name__ == "__main__":
+    unittest.main()