python-sdk/tests/server/lowlevel/test_handler_registration.py at c413473d42a3dc7122207cdd49ca3e086f15cb8f · modelcontextprotocol/python-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
"""Tests for public handler registration/deregistration API on low-level Server."""

from typing import Any

import pytest

from mcp.server.context import ServerRequestContext
from mcp.server.lowlevel.server import Server


@pytest.fixture
def server() -> Server[Any]:
    return Server(name="test-server")


async def _dummy_request_handler(ctx: ServerRequestContext[Any], params: Any) -> dict[str, str]:  # pragma: no cover
    return {"result": "ok"}


async def _dummy_notification_handler(ctx: ServerRequestContext[Any], params: Any) -> None:  # pragma: no cover
    pass


class TestAddRequestHandler:
    def test_add_request_handler(self, server: Server[Any]) -> None:
        server.add_request_handler("custom/method", _dummy_request_handler)
        assert server.has_handler("custom/method")

    def test_add_request_handler_replaces_existing(self, server: Server[Any]) -> None:
        async def handler_a(ctx: ServerRequestContext[Any], params: Any) -> str:  # pragma: no cover
            return "a"

        async def handler_b(ctx: ServerRequestContext[Any], params: Any) -> str:  # pragma: no cover
            return "b"

        server.add_request_handler("custom/method", handler_a)
        server.add_request_handler("custom/method", handler_b)
        # The second handler should replace the first
        assert server._request_handlers["custom/method"] is handler_b


class TestRemoveRequestHandler:
    def test_remove_request_handler(self, server: Server[Any]) -> None:
        server.add_request_handler("custom/method", _dummy_request_handler)
        assert server.has_handler("custom/method")
        server.remove_request_handler("custom/method")
        assert not server.has_handler("custom/method")

    def test_remove_request_handler_not_found(self, server: Server[Any]) -> None:
        with pytest.raises(KeyError):
            server.remove_request_handler("nonexistent/method")


class TestAddNotificationHandler:
    def test_add_notification_handler(self, server: Server[Any]) -> None:
        server.add_notification_handler("custom/notify", _dummy_notification_handler)
        assert server.has_handler("custom/notify")

    def test_add_notification_handler_replaces_existing(self, server: Server[Any]) -> None:
        async def handler_a(ctx: ServerRequestContext[Any], params: Any) -> None:  # pragma: no cover
            pass

        async def handler_b(ctx: ServerRequestContext[Any], params: Any) -> None:  # pragma: no cover
            pass

        server.add_notification_handler("custom/notify", handler_a)
        server.add_notification_handler("custom/notify", handler_b)
        assert server._notification_handlers["custom/notify"] is handler_b


class TestRemoveNotificationHandler:
    def test_remove_notification_handler(self, server: Server[Any]) -> None:
        server.add_notification_handler("custom/notify", _dummy_notification_handler)
        assert server.has_handler("custom/notify")
        server.remove_notification_handler("custom/notify")
        assert not server.has_handler("custom/notify")

    def test_remove_notification_handler_not_found(self, server: Server[Any]) -> None:
        with pytest.raises(KeyError):
            server.remove_notification_handler("nonexistent/notify")


class TestHasHandler:
    def test_has_handler_request(self, server: Server[Any]) -> None:
        server.add_request_handler("custom/method", _dummy_request_handler)
        assert server.has_handler("custom/method")

    def test_has_handler_notification(self, server: Server[Any]) -> None:
        server.add_notification_handler("custom/notify", _dummy_notification_handler)
        assert server.has_handler("custom/notify")

    def test_has_handler_unregistered(self, server: Server[Any]) -> None:
        assert not server.has_handler("nonexistent/method")

    def test_has_handler_default_ping(self, server: Server[Any]) -> None:
        """The ping handler is registered by default."""
        assert server.has_handler("ping")