from unittest import TestCase
from unittest.mock import patch

from sglang.srt.managers.scheduler_dp_attn_mixin import (
    MLPSyncBatchInfo,
    prepare_mlp_sync_batch_raw,
)


class _FakeTPGroup:
    device_group = object()
    cpu_group = object()
    device = "cuda"


class TestSchedulerDPAttnMixin(TestCase):
    def test_single_dp_idle_batch_does_not_enter_mlp_sync_collective(self):
        def fail_all_gather(self, *args, **kwargs):
            raise AssertionError("idle single-DP scheduler must not all-gather")

        with patch.object(MLPSyncBatchInfo, "all_gather", fail_all_gather):
            result = prepare_mlp_sync_batch_raw(
                local_batch=None,
                dp_size=1,
                attn_tp_size=1,
                attn_cp_size=8,
                tp_group=_FakeTPGroup(),
                get_idle_batch=lambda: (_ for _ in ()).throw(
                    AssertionError("idle single-DP scheduler must not build idle batch")
                ),
                disable_cuda_graph=False,
                require_mlp_tp_gather=True,
                disable_overlap_schedule=True,
                offload_tags=set(),
            )

        self.assertIsNone(result)