# Copyright (c) OpenMMLab. All rights reserved.
import pytest

from lmdeploy.messages import PytorchEngineConfig
from lmdeploy.pytorch.disagg.config import EngineRole
from lmdeploy.pytorch.engine.engine_checker import EngineChecker


def _build_checker(engine_config: PytorchEngineConfig):
    checker = object.__new__(EngineChecker)
    checker.engine_config = engine_config

    def log_and_raise(*args, **kwargs):
        raise RuntimeError(kwargs.get('message'))

    checker.log_and_exit = log_and_raise
    return checker


def test_engine_checker_rejects_split_kernel_blocks_for_pd_migration():
    engine_config = PytorchEngineConfig(max_batch_size=1,
                                        role=EngineRole.Prefill,
                                        block_size=64,
                                        kernel_block_size=32)
    checker = _build_checker(engine_config)

    with pytest.raises(RuntimeError, match='PD migration does not support block_size != kernel_block_size'):
        checker.check()


def test_engine_checker_allows_split_kernel_blocks_for_hybrid_engine():
    engine_config = PytorchEngineConfig(max_batch_size=1,
                                        role=EngineRole.Hybrid,
                                        block_size=64,
                                        kernel_block_size=32)
    checker = _build_checker(engine_config)

    checker.check()