/test/registered/unit/plugins/test_hook_registry.py

"""
Unit tests for the hook registry system.

Covers: basic hooks (AROUND/BEFORE/AFTER/REPLACE), descriptor preservation
(classmethod/staticmethod), hook ordering, cross-target conflict detection,
patch propagation, and edge cases.

Run:  python -m pytest test/registered/unit/plugins/test_hook_registry.py -v
"""

import sys
import types
import uuid

from sglang.srt.plugins.hook_registry import HookRegistry, HookType, plugin_hook
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

register_cpu_ci(est_time=7, suite="base-a-test-cpu")

# ---------------------------------------------------------------------------
# Helpers: synthetic module creation
# ---------------------------------------------------------------------------

_SYNTH_MODULE_PREFIX = "_synth_hook_test_"


def _make_module(**attrs):
    """Create a throwaway module registered in sys.modules."""
    name = f"{_SYNTH_MODULE_PREFIX}{uuid.uuid4().hex[:8]}"
    mod = types.ModuleType(name)
    for k, v in attrs.items():
        setattr(mod, k, v)
    sys.modules[name] = mod
    return mod, name


def _cleanup_synth_modules():
    """Remove all synthetic modules from sys.modules."""
    to_del = [k for k in sys.modules if k.startswith(_SYNTH_MODULE_PREFIX)]
    for k in to_del:
        del sys.modules[k]


# ---------------------------------------------------------------------------
# Base class for hook tests (shared setUp/tearDown)
# ---------------------------------------------------------------------------


class _HookTestCase(CustomTestCase):
    """Base class that resets HookRegistry and cleans up synth modules."""

    def setUp(self):
        HookRegistry.reset()
        _cleanup_synth_modules()

    def tearDown(self):
        HookRegistry.reset()
        _cleanup_synth_modules()


# ===========================================================================
# TestBasicHooks
# ===========================================================================


class TestBasicHooks(_HookTestCase):
    """AROUND / BEFORE / AFTER / REPLACE on plain functions, class REPLACE,
    and the @plugin_hook decorator."""

    def test_around_function(self):
        def orig(x):
            return x * 2

        mod, name = _make_module(orig=orig)

        def add_one(original_fn, x):
            return original_fn(x) + 1

        HookRegistry.register(f"{name}.orig", add_one, HookType.AROUND)
        HookRegistry.apply_hooks()
        self.assertEqual(mod.orig(3), 7)  # 3*2 + 1

    def test_before_modifies_args(self):
        """BEFORE hook returns (args, kwargs) to modify arguments."""

        def orig(x, y=0):
            return x + y

        mod, name = _make_module(orig=orig)

        def double_x(x, y=0):
            return (x * 2,), {"y": y + 1}

        HookRegistry.register(f"{name}.orig", double_x, HookType.BEFORE)
        HookRegistry.apply_hooks()
        self.assertEqual(mod.orig(3), 7)  # x=3*2=6, y=0+1=1, 6+1=7

    def test_before_returning_none(self):
        """BEFORE hook returning None leaves arguments unchanged."""

        def orig(x):
            return x * 2

        mod, name = _make_module(orig=orig)

        def before_noop(x):
            return None  # leave args unchanged

        HookRegistry.register(f"{name}.orig", before_noop, HookType.BEFORE)
        HookRegistry.apply_hooks()
        self.assertEqual(mod.orig(3), 6)  # args unchanged

    def test_after_function(self):
        def orig(x):
            return x * 2

        mod, name = _make_module(orig=orig)

        def add_ten(result, x):
            return result + 10

        HookRegistry.register(f"{name}.orig", add_ten, HookType.AFTER)
        HookRegistry.apply_hooks()
        self.assertEqual(mod.orig(3), 16)  # 3*2 + 10

    def test_replace_function(self):
        def orig(x):
            return x * 2

        mod, name = _make_module(orig=orig)

        def replacement(x):
            return x * 100

        HookRegistry.register(f"{name}.orig", replacement, HookType.REPLACE)
        HookRegistry.apply_hooks()
        self.assertEqual(mod.orig(3), 300)

    def test_class_replace(self):
        class Original:
            def greet(self):
                return "original"

        mod, name = _make_module(Original=Original)

        class Replacement(Original):
            def greet(self):
                return "replaced"

        HookRegistry.register(f"{name}.Original", Replacement, HookType.REPLACE)
        HookRegistry.apply_hooks()

        self.assertIs(mod.Original, Replacement)
        self.assertIsInstance(mod.Original(), Replacement)
        self.assertEqual(mod.Original().greet(), "replaced")

    def test_plugin_hook_decorator(self):
        def orig(x):
            return x

        mod, name = _make_module(orig=orig)

        @plugin_hook(f"{name}.orig", type=HookType.REPLACE)
        def my_replace(x):
            return x + 42

        HookRegistry.apply_hooks()
        self.assertEqual(mod.orig(0), 42)


# ===========================================================================
# TestDescriptorPreservation  (Bug B regression tests)
# ===========================================================================


class TestDescriptorPreservation(_HookTestCase):
    """Hooks on classmethod/staticmethod must preserve descriptor semantics."""

    def _make_cls_module(self):
        class MyClass:
            @classmethod
            def cm(cls, x):
                return ("cm", cls.__name__, x)

            @staticmethod
            def sm(x):
                return ("sm", x)

        mod, name = _make_module(MyClass=MyClass)
        return mod, name, MyClass

    def test_around_classmethod(self):
        mod, name, MyClass = self._make_cls_module()

        def add_tag(original_fn, cls, x):
            return original_fn(cls, x) + ("around",)

        HookRegistry.register(f"{name}.MyClass.cm", add_tag, HookType.AROUND)
        HookRegistry.apply_hooks()

        result = mod.MyClass.cm(1)
        self.assertEqual(result, ("cm", "MyClass", 1, "around"))

    def test_replace_classmethod(self):
        mod, name, MyClass = self._make_cls_module()

        def new_cm(cls, x):
            return ("replaced_cm", cls.__name__, x)

        HookRegistry.register(f"{name}.MyClass.cm", new_cm, HookType.REPLACE)
        HookRegistry.apply_hooks()

        result = mod.MyClass.cm(1)
        self.assertEqual(result, ("replaced_cm", "MyClass", 1))

    def test_around_staticmethod(self):
        mod, name, MyClass = self._make_cls_module()

        def wrap_sm(original_fn, x):
            return original_fn(x) + ("around",)

        HookRegistry.register(f"{name}.MyClass.sm", wrap_sm, HookType.AROUND)
        HookRegistry.apply_hooks()

        result = mod.MyClass.sm(1)
        self.assertEqual(result, ("sm", 1, "around"))

    def test_replace_staticmethod(self):
        mod, name, MyClass = self._make_cls_module()

        def new_sm(x):
            return ("replaced_sm", x)

        HookRegistry.register(f"{name}.MyClass.sm", new_sm, HookType.REPLACE)
        HookRegistry.apply_hooks()

        result = mod.MyClass.sm(1)
        self.assertEqual(result, ("replaced_sm", 1))

    def test_classmethod_subclass_cls(self):
        mod, name, MyClass = self._make_cls_module()

        def add_tag(original_fn, cls, x):
            return original_fn(cls, x) + ("around",)

        HookRegistry.register(f"{name}.MyClass.cm", add_tag, HookType.AROUND)
        HookRegistry.apply_hooks()

        class Sub(mod.MyClass):
            pass

        result = Sub.cm(1)
        self.assertEqual(result, ("cm", "Sub", 1, "around"))


# ===========================================================================
# TestHookOrdering
# ===========================================================================


class TestHookOrdering(_HookTestCase):
    """Verify REPLACE is applied first, then other hooks wrap it."""

    def test_replace_then_around(self):
        def orig(x):
            return x

        mod, name = _make_module(orig=orig)

        def repl(x):
            return x * 10

        def add_one(original_fn, x):
            return original_fn(x) + 1

        HookRegistry.register(f"{name}.orig", repl, HookType.REPLACE)
        HookRegistry.register(f"{name}.orig", add_one, HookType.AROUND)
        HookRegistry.apply_hooks()
        # REPLACE first: x*10, then AROUND: +1  => 31
        self.assertEqual(mod.orig(3), 31)

    def test_replace_before_after(self):
        def orig(x):
            return x

        mod, name = _make_module(orig=orig)

        def repl(x):
            return x * 10

        def double_arg(x):
            return (x * 2,), {}

        def add_hundred(result, x):
            return result + 100

        HookRegistry.register(f"{name}.orig", repl, HookType.REPLACE)
        HookRegistry.register(f"{name}.orig", double_arg, HookType.BEFORE)
        HookRegistry.register(f"{name}.orig", add_hundred, HookType.AFTER)
        HookRegistry.apply_hooks()
        # BEFORE doubles x: 3*2=6 → REPLACE: 6*10=60 → AFTER: 60+100=160
        self.assertEqual(mod.orig(3), 160)


# ===========================================================================
# TestCrossTargetConflict
# ===========================================================================


class TestCrossTargetConflict(_HookTestCase):
    """Verify warning for class REPLACE + method REPLACE combo."""

    def test_class_replace_then_method_replace_warns(self):
        class Original:
            def foo(self):
                return "orig"

        mod, name = _make_module(Original=Original)

        class Replacement(Original):
            def foo(self):
                return "class_replaced"

        HookRegistry.register(f"{name}.Original", Replacement, HookType.REPLACE)

        def method_repl(self):
            return "method_replaced"

        HookRegistry.register(f"{name}.Original.foo", method_repl, HookType.REPLACE)

        with self.assertLogs("sglang.srt.plugins.hook_registry", level="WARNING") as cm:
            HookRegistry.apply_hooks()

        self.assertTrue(any("will override" in msg for msg in cm.output))


# ===========================================================================
# TestPatchPropagation
# ===========================================================================


class TestPatchPropagation(_HookTestCase):
    """Verify that patches propagate to other modules that imported the target."""

    def test_same_reference_propagates(self):
        def orig(x):
            return x * 2

        source_mod, source_name = _make_module(orig=orig)
        importer_mod, _ = _make_module(orig=orig)  # same reference

        def add_one(fn, x):
            return fn(x) + 1

        HookRegistry.register(f"{source_name}.orig", add_one, HookType.AROUND)
        HookRegistry.apply_hooks()

        self.assertEqual(source_mod.orig(3), 7)
        self.assertEqual(importer_mod.orig(3), 7)


# ===========================================================================
# TestEdgeCases
# ===========================================================================


class TestEdgeCases(_HookTestCase):
    """Reset, type validation, multi-AROUND onion, idempotent apply."""

    def test_reset(self):
        def orig(x):
            return x

        mod, name = _make_module(orig=orig)

        def noop(fn, x):
            return fn(x)

        HookRegistry.register(f"{name}.orig", noop, HookType.AROUND)
        HookRegistry.reset()

        HookRegistry.apply_hooks()
        self.assertEqual(mod.orig(3), 3)

    def test_register_class_with_wrong_type(self):
        class BadHook:
            pass

        for ht in (HookType.BEFORE, HookType.AFTER, HookType.AROUND):
            with self.assertRaises(TypeError):
                HookRegistry.register("some.target", BadHook, ht)

    def test_multi_around_onion(self):
        call_order = []

        def orig(x):
            call_order.append("orig")
            return x

        mod, name = _make_module(orig=orig)

        def around1(fn, x):
            call_order.append("a1_before")
            result = fn(x)
            call_order.append("a1_after")
            return result + 1

        def around2(fn, x):
            call_order.append("a2_before")
            result = fn(x)
            call_order.append("a2_after")
            return result + 10

        HookRegistry.register(f"{name}.orig", around1, HookType.AROUND)
        HookRegistry.register(f"{name}.orig", around2, HookType.AROUND)
        HookRegistry.apply_hooks()

        result = mod.orig(0)
        self.assertEqual(result, 11)
        self.assertEqual(
            call_order, ["a2_before", "a1_before", "orig", "a1_after", "a2_after"]
        )

    def test_apply_idempotent(self):
        call_count = [0]

        def orig(x):
            return x

        mod, name = _make_module(orig=orig)

        def counter(fn, x):
            call_count[0] += 1
            return fn(x)

        HookRegistry.register(f"{name}.orig", counter, HookType.AROUND)
        HookRegistry.apply_hooks()
        HookRegistry.apply_hooks()  # second apply should be no-op

        mod.orig(1)
        self.assertEqual(call_count[0], 1)


if __name__ == "__main__":
    import unittest

    unittest.main()