"""Unit tests for the async/non-blocking mooncake wrapper bindings (G1).

batch_transfer_async_submit() must do a non-blocking submit returning a batch_id;
wait_batch_transfers() must block-until-all and return 0/non-zero. Both must raise a
clear upgrade error if the engine lacks the method, and degrade to -1 on transient
engine errors. No GPU / real RDMA needed — the engine is mocked.
"""

import unittest

from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
    MooncakeTransferEngine,
)


class _GoodEngine:
    def __init__(self, status_ret=0):
        self.calls = []
        self._status_ret = status_ret

    def batch_transfer_async_write(self, sid, buffers, peers, lengths):
        self.calls.append(("submit", sid, list(buffers), list(peers), list(lengths)))
        return 7  # batch_id handle

    def get_batch_transfer_status(self, ids):
        self.calls.append(("wait", list(ids)))
        return self._status_ret


class _NoAsyncEngine:
    pass  # lacks both methods


class _RaisingEngine:
    def batch_transfer_async_write(self, *a):
        raise RuntimeError("transient rdma error")

    def get_batch_transfer_status(self, *a):
        raise RuntimeError("transient rdma error")


def _wrapper(engine):
    w = MooncakeTransferEngine.__new__(MooncakeTransferEngine)  # skip real init/connect
    w.engine = engine
    return w


class TestAsyncSubmit(unittest.TestCase):
    def test_submit_returns_batch_id_and_dispatches(self):
        eng = _GoodEngine()
        w = _wrapper(eng)
        bid = w.batch_transfer_async_submit("sess", [1, 2], [10, 20], [64, 64])
        self.assertEqual(bid, 7)
        self.assertEqual(eng.calls[0], ("submit", "sess", [1, 2], [10, 20], [64, 64]))

    def test_submit_missing_method_raises_upgrade(self):
        w = _wrapper(_NoAsyncEngine())
        with self.assertRaises(RuntimeError):
            w.batch_transfer_async_submit("s", [1], [2], [3])

    def test_submit_transient_error_returns_minus_one(self):
        w = _wrapper(_RaisingEngine())
        self.assertEqual(w.batch_transfer_async_submit("s", [1], [2], [3]), -1)


class TestWait(unittest.TestCase):
    def test_wait_dispatches_and_returns_status(self):
        eng = _GoodEngine(status_ret=0)
        w = _wrapper(eng)
        self.assertEqual(w.wait_batch_transfers([7, 8]), 0)
        self.assertEqual(eng.calls[0], ("wait", [7, 8]))

    def test_wait_empty_is_zero_without_engine_call(self):
        eng = _GoodEngine()
        w = _wrapper(eng)
        self.assertEqual(w.wait_batch_transfers([]), 0)
        self.assertEqual(eng.calls, [])

    def test_wait_nonzero_status_propagates(self):
        w = _wrapper(_GoodEngine(status_ret=-1))
        self.assertEqual(w.wait_batch_transfers([7]), -1)

    def test_wait_missing_method_raises_upgrade(self):
        w = _wrapper(_NoAsyncEngine())
        with self.assertRaises(RuntimeError):
            w.wait_batch_transfers([7])

    def test_wait_transient_error_returns_minus_one(self):
        w = _wrapper(_RaisingEngine())
        self.assertEqual(w.wait_batch_transfers([7]), -1)


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