test_simple_engine_unit.py

# SPDX-License-Identifier: Apache-2.0
"""
Tests for SimpleEngine and GenerationOutput.

Generated by MiniMax-M2.5 via openclaw simulation, reviewed and fixed by Claude.
"""

import asyncio
from unittest.mock import MagicMock, patch

import pytest

from vllm_mlx.engine.base import GenerationOutput
from vllm_mlx.engine.simple import SimpleEngine

# ---------------------------------------------------------------------------
# GenerationOutput dataclass
# ---------------------------------------------------------------------------


class TestGenerationOutput:
    """Tests for GenerationOutput dataclass defaults."""

    def test_default_values(self):
        output = GenerationOutput(text="test")
        assert output.text == "test"
        assert output.tokens == []
        assert output.prompt_tokens == 0
        assert output.completion_tokens == 0
        assert output.finish_reason == "stop"
        assert output.new_text == ""
        assert output.finished is True
        assert output.logprobs is None

    def test_all_values_provided(self):
        output = GenerationOutput(
            text="generated",
            tokens=[1, 2, 3],
            prompt_tokens=10,
            completion_tokens=5,
            finish_reason="length",
            new_text=" new",
            finished=False,
            logprobs="mock",
        )
        assert output.tokens == [1, 2, 3]
        assert output.finish_reason == "length"
        assert output.finished is False
        assert output.logprobs == "mock"

    def test_none_finish_reason(self):
        output = GenerationOutput(text="x", finish_reason=None)
        assert output.finish_reason is None


# ---------------------------------------------------------------------------
# SimpleEngine initialization
# ---------------------------------------------------------------------------


class TestSimpleEngineInit:
    """Tests for SimpleEngine initialization."""

    def test_default_params(self):
        engine = SimpleEngine(model_name="test-model")
        assert engine._model_name == "test-model"
        assert engine._trust_remote_code is True
        assert engine._enable_cache is True
        assert engine._draft_model_name is None
        assert engine._num_draft_tokens == 4
        assert engine._prefill_step_size == 2048
        assert engine._kv_bits is None
        assert engine._kv_group_size == 64
        assert engine._model is None
        assert engine._loaded is False

    def test_custom_params(self):
        engine = SimpleEngine(
            model_name="custom",
            trust_remote_code=False,
            enable_cache=False,
            force_mllm=True,
            draft_model="draft",
            num_draft_tokens=8,
            prefill_step_size=1024,
            kv_bits=4,
            kv_group_size=32,
        )
        assert engine._trust_remote_code is False
        assert engine._is_mllm is True
        assert engine._draft_model_name == "draft"
        assert engine._prefill_step_size == 1024
        assert engine._kv_bits == 4

    @patch("vllm_mlx.engine.simple.is_mllm_model")
    def test_mllm_auto_detection(self, mock_is_mllm):
        mock_is_mllm.return_value = True
        engine = SimpleEngine(model_name="vision-model")
        assert engine._is_mllm is True

    @patch("vllm_mlx.engine.simple.is_mllm_model")
    def test_force_mllm_overrides_detection(self, mock_is_mllm):
        mock_is_mllm.return_value = False
        engine = SimpleEngine(model_name="text-only", force_mllm=True)
        assert engine._is_mllm is True


# ---------------------------------------------------------------------------
# Properties
# ---------------------------------------------------------------------------


class TestSimpleEngineProperties:
    """Tests for SimpleEngine properties."""

    def test_model_name(self):
        assert SimpleEngine(model_name="foo").model_name == "foo"

    def test_is_mllm(self):
        assert SimpleEngine(model_name="x").is_mllm is False
        assert SimpleEngine(model_name="x", force_mllm=True).is_mllm is True

    def test_tokenizer_before_load(self):
        assert SimpleEngine(model_name="x").tokenizer is None

    def test_tokenizer_llm_after_load(self):
        engine = SimpleEngine(model_name="x")
        mock_tok = MagicMock()
        mock_model = MagicMock()
        mock_model.tokenizer = mock_tok
        engine._model = mock_model
        engine._loaded = True
        assert engine.tokenizer is mock_tok

    def test_tokenizer_mllm_after_load(self):
        engine = SimpleEngine(model_name="x", force_mllm=True)
        mock_proc = MagicMock()
        mock_model = MagicMock()
        mock_model.processor = mock_proc
        engine._model = mock_model
        engine._loaded = True
        assert engine.tokenizer is mock_proc


# ---------------------------------------------------------------------------
# Start / Stop
# ---------------------------------------------------------------------------


class TestStartStop:
    """Tests for start and stop methods."""

    def test_start_llm(self):
        engine = SimpleEngine(model_name="test")
        with patch("vllm_mlx.models.llm.MLXLanguageModel") as MockLLM:
            mock_inst = MagicMock()
            MockLLM.return_value = mock_inst
            asyncio.run(engine.start())
            mock_inst.load.assert_called_once()
            assert engine._loaded is True

    def test_start_idempotent(self):
        async def _run():
            engine = SimpleEngine(model_name="test")
            with patch("vllm_mlx.models.llm.MLXLanguageModel") as MockLLM:
                MockLLM.return_value = MagicMock()
                await engine.start()
                await engine.start()
                MockLLM.assert_called_once()

        asyncio.run(_run())

    def test_stop_clears_state(self):
        engine = SimpleEngine(model_name="test")
        engine._model = MagicMock()
        engine._loaded = True
        asyncio.run(engine.stop())
        assert engine._model is None
        assert engine._loaded is False


# ---------------------------------------------------------------------------
# stream_generate
# ---------------------------------------------------------------------------


def _make_chunk(text, finished=False, finish_reason=None, prompt_tokens=5, token=1):
    """Helper to create mock streaming chunks."""
    c = MagicMock(spec=[])  # spec=[] prevents auto-creating attributes
    c.text = text
    c.finished = finished
    c.finish_reason = finish_reason
    c.prompt_tokens = prompt_tokens
    c.token = token
    c.logprobs = None
    return c


def _collect_stream(engine, **kwargs):
    """Helper to collect async generator results synchronously."""

    async def _run():
        outputs = []
        async for out in engine.stream_generate(**kwargs):
            outputs.append(out)
        return outputs

    return asyncio.run(_run())


class TestStreamGenerate:
    """Tests for stream_generate method."""

    @pytest.fixture
    def engine(self):
        e = SimpleEngine(model_name="test")
        e._model = MagicMock()
        e._loaded = True
        return e

    def test_accumulates_text(self, engine):
        chunks = [
            _make_chunk("Hello", token=1),
            _make_chunk(" World", finished=True, finish_reason="stop", token=2),
        ]
        engine._model.stream_generate = MagicMock(return_value=iter(chunks))

        outputs = _collect_stream(engine, prompt="test", max_tokens=100)

        assert len(outputs) == 2
        assert outputs[0].text == "Hello"
        assert outputs[0].new_text == "Hello"
        assert outputs[0].finished is False
        assert outputs[1].text == "Hello World"
        assert outputs[1].new_text == " World"
        assert outputs[1].finished is True
        assert outputs[1].finish_reason == "stop"

    def test_max_tokens_finish(self, engine):
        """When max_tokens reached, generation stops."""
        chunks = [
            _make_chunk("a", token=1),
            _make_chunk("b", token=2),
            _make_chunk("c", token=3),
            _make_chunk("d", token=4),
            _make_chunk("e", token=5),
        ]
        engine._model.stream_generate = MagicMock(return_value=iter(chunks))

        outputs = _collect_stream(engine, prompt="test", max_tokens=3)

        assert outputs[-1].finished is True
        assert outputs[-1].completion_tokens <= 4

    def test_prompt_tokens_propagated(self, engine):
        chunks = [
            _make_chunk("x", prompt_tokens=42, finished=True, finish_reason="stop"),
        ]
        engine._model.stream_generate = MagicMock(return_value=iter(chunks))

        outputs = _collect_stream(engine, prompt="test", max_tokens=100)
        assert outputs[0].prompt_tokens == 42

    def test_empty_stream_yields_final(self, engine):
        """If model yields no chunks at all, still get a finish output."""
        engine._model.stream_generate = MagicMock(return_value=iter([]))

        outputs = _collect_stream(engine, prompt="test", max_tokens=100)

        assert len(outputs) == 1
        assert outputs[0].finished is True
        assert outputs[0].text == ""

    def test_logprobs_passthrough(self, engine):
        chunk = _make_chunk("x", finished=True, finish_reason="stop")
        chunk.logprobs = "mock_logprobs"
        engine._model.stream_generate = MagicMock(return_value=iter([chunk]))

        outputs = _collect_stream(engine, prompt="test", max_tokens=100)
        assert outputs[0].logprobs == "mock_logprobs"


# ---------------------------------------------------------------------------
# get_stats
# ---------------------------------------------------------------------------


class TestGetStats:
    """Tests for get_stats method."""

    def test_basic_stats(self):
        engine = SimpleEngine(model_name="test-model")
        stats = engine.get_stats()
        assert stats["engine_type"] == "simple"
        assert stats["model_name"] == "test-model"
        assert stats["is_mllm"] is False
        assert stats["loaded"] is False

    def test_stats_after_load(self):
        engine = SimpleEngine(model_name="test-model")
        engine._loaded = True
        assert engine.get_stats()["loaded"] is True


# ---------------------------------------------------------------------------
# preserve_native_tool_format
# ---------------------------------------------------------------------------


class TestPreserveNativeToolFormat:
    """Tests for preserve_native_tool_format property."""

    def test_default_false(self):
        engine = SimpleEngine(model_name="test")
        assert engine.preserve_native_tool_format is False

    def test_set_and_get(self):
        engine = SimpleEngine(model_name="test")
        engine.preserve_native_tool_format = True
        assert engine.preserve_native_tool_format is True

    def test_set_false(self):
        engine = SimpleEngine(model_name="test")
        engine.preserve_native_tool_format = True
        engine.preserve_native_tool_format = False
        assert engine.preserve_native_tool_format is False


# ---------------------------------------------------------------------------
# _inject_shared_model config propagation
# ---------------------------------------------------------------------------


class TestInjectSharedModelConfig:
    """Tests for _inject_shared_model using engine config values."""

    @pytest.mark.asyncio
    async def test_inject_propagates_engine_config(self):
        """Injected model uses engine's config, not hardcoded defaults (regression)."""
        engine = SimpleEngine(
            model_name="test",
            prefill_step_size=4096,
            kv_bits=4,
            kv_group_size=128,
        )

        mock_model = MagicMock()
        mock_tokenizer = MagicMock()

        await engine._inject_shared_model(mock_model, mock_tokenizer)

        assert engine._model.prefill_step_size == 4096
        assert engine._model.kv_bits == 4
        assert engine._model.kv_group_size == 128

    @pytest.mark.asyncio
    async def test_inject_default_config(self):
        """Injected model uses default config when engine uses defaults."""
        engine = SimpleEngine(model_name="test")

        mock_model = MagicMock()
        mock_tokenizer = MagicMock()

        await engine._inject_shared_model(mock_model, mock_tokenizer)

        assert engine._model.prefill_step_size == 2048
        assert engine._model.kv_bits is None
        assert engine._model.kv_group_size == 64