import importlib.metadata
import importlib.util
import os
import unittest
from unittest import mock

from tests.helpers import TmpCase


@unittest.skipIf(importlib.util.find_spec("vllm") is None, "needs vllm")
class Plugin(TmpCase):
    def test_entry_point_is_installed(self):
        plugins = importlib.metadata.entry_points(group="vllm.general_plugins")
        self.assertIn("prain", [p.name for p in plugins])

    def test_register_swaps_in_tiered_experts(self):
        config = self.write("p.toml", f'pack_dir = "{self.dir}/pack"\n')
        with mock.patch.dict(os.environ, {"PRAIN_CONFIG": config}):
            from prain.plugin import register

            register()
            register()  # vLLM may call it again; must not raise
        from vllm.model_executor.custom_op import op_registry_oot

        from prain.experts import TieredRoutedExperts

        self.assertIs(op_registry_oot["RoutedExperts"], TieredRoutedExperts)

    def test_draft_hook_wraps_the_v2_model_runner_once(self):
        from vllm.v1.worker.gpu.input_batch import InputBatch
        from vllm.v1.worker.gpu.model_runner import GPUModelRunner

        from prain.experts import _install_draft_hook

        _install_draft_hook()
        wrapped = GPUModelRunner.prepare_inputs
        _install_draft_hook()
        self.assertIs(GPUModelRunner.prepare_inputs, wrapped)
        self.assertTrue(wrapped.prain)
        self.assertLessEqual({"req_ids", "query_start_loc_np"}, set(InputBatch.__dataclass_fields__))