Improve test coverage

Signed-off-by: Mihai Criveti <[email protected]>

Author: crivetimihai

tests/unit/mcpgateway/cache/test_session_registry.py

@@ -0,0 +1,119 @@
+# -*- coding: utf-8 -*-
+"""Memory-backend unit tests for `session_registry.py`.
+
+Copyright 2025
+SPDX-License-Identifier: Apache-2.0
+Author: Mihai Criveti
+
+These tests cover the essential public behaviours of SessionRegistry when
+configured with backend="memory":
+
+* add_session / get_session / get_session_sync
+* remove_session (disconnects transport & clears cache)
+* broadcast + respond (with generate_response monkey-patched)
+
+No Redis or SQLAlchemy fixtures are required, making the suite fast and
+portable.
+"""
+
+import asyncio
+
+import pytest
+
+# Import SessionRegistry – works whether the file lives inside the package or beside it
+try:
+    from mcpgateway.cache.session_registry import SessionRegistry
+except (ModuleNotFoundError, ImportError):  # pragma: no cover
+    from session_registry import SessionRegistry  # type: ignore
+
+
+class FakeSSETransport:
+    """Minimal stub implementing only the methods SessionRegistry uses."""
+
+    def __init__(self, session_id: str):
+        self.session_id = session_id
+        self._connected = True
+        self.sent_messages = []
+
+    async def disconnect(self):
+        self._connected = False
+
+    async def is_connected(self):
+        return self._connected
+
+    async def send_message(self, msg):
+        self.sent_messages.append(msg)
+
+
+# ---------------------------------------------------------------------------
+# Pytest fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(name="event_loop")
+def _event_loop_fixture():
+    loop = asyncio.new_event_loop()
+    yield loop
+    loop.close()
+
+
+@pytest.fixture()
+async def registry():
+    reg = SessionRegistry(backend="memory")
+    await reg.initialize()
+    yield reg
+    await reg.shutdown()
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_add_and_get_session(registry):
+    tr = FakeSSETransport("abc")
+    await registry.add_session("abc", tr)
+
+    assert await registry.get_session("abc") is tr
+    assert registry.get_session_sync("abc") is tr
+
+
+@pytest.mark.asyncio
+async def test_remove_session(registry):
+    tr = FakeSSETransport("dead")
+    await registry.add_session("dead", tr)
+
+    await registry.remove_session("dead")
+
+    assert not await tr.is_connected()
+    assert registry.get_session_sync("dead") is None
+
+
+@pytest.mark.asyncio
+async def test_broadcast_and_respond(monkeypatch, registry):
+    """Ensure broadcast stores the message and respond delivers it via generate_response."""
+    tr = FakeSSETransport("xyz")
+    await registry.add_session("xyz", tr)
+
+    captured = {}
+
+    async def fake_generate_response(*, message, transport, **_):
+        captured["transport"] = transport
+        captured["message"] = message
+
+    monkeypatch.setattr(registry, "generate_response", fake_generate_response)
+
+    ping_msg = {"method": "ping", "id": 1, "params": {}}
+    await registry.broadcast("xyz", ping_msg)
+
+    # respond should call our fake_generate_response exactly once
+    await registry.respond(
+        server_id=None,
+        user={},
+        session_id="xyz",
+        base_url="http://localhost",
+    )
+
+    assert captured["transport"] is tr
+    assert captured["message"] == ping_msg

tests/unit/mcpgateway/test_discovery.py

@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+"""Unit tests for Federation Discovery Service.
+
+Copyright 2025
+SPDX-License-Identifier: Apache-2.0
+Authors: Mihai Criveti
+
+Comprehensive unit tests for the discovery service module.
+"""
+
+# tests/test_discovery.py
+import asyncio
+from datetime import datetime
+
+import pytest
+
+from mcpgateway.federation.discovery import (
+    PROTOCOL_VERSION,
+    DiscoveryService,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+async def discovery():
+    """
+    Provide a DiscoveryService instance whose network-touching method
+    `_get_gateway_info` is monkey-patched so no real HTTP requests are made.
+    """
+    ds = DiscoveryService()
+
+    async def _fake_gateway_info(url: str):  # noqa: D401, ANN001
+        # Return an *empty* capabilities object – structure is unimportant here.
+        from mcpgateway.types import ServerCapabilities
+
+        return ServerCapabilities()
+
+    # Patch the network call
+    ds._get_gateway_info = _fake_gateway_info  # type: ignore[attr-defined]
+
+    yield ds
+    await ds.stop()  # ensure graceful cleanup
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.anyio
+async def test_add_peer_success(discovery):
+    added = await discovery.add_peer("http://example.com", source="test", name="Example")
+    assert added is True, "first call should add the peer"
+
+    peers = discovery.get_discovered_peers()
+    assert len(peers) == 1
+    peer = peers[0]
+
+    assert peer.url == "http://example.com"
+    assert peer.name == "Example"
+    assert peer.protocol_version == PROTOCOL_VERSION
+    assert peer.source == "test"
+
+
+@pytest.mark.anyio
+async def test_add_duplicate_peer_is_ignored(discovery):
+    await discovery.add_peer("http://dup.com", source="test")
+    second_add = await discovery.add_peer("http://dup.com", source="test-again")
+    assert second_add is False, "duplicate add should be a no-op"
+
+    peers = discovery.get_discovered_peers()
+    assert len(peers) == 1, "still only one peer stored"
+
+
+@pytest.mark.anyio
+async def test_add_peer_invalid_url_returns_false(discovery):
+    added = await discovery.add_peer("not-a-valid-url", source="test")
+    assert added is False
+    assert discovery.get_discovered_peers() == []
+
+
+@pytest.mark.anyio
+async def test_refresh_peer_updates_last_seen(discovery):
+    await discovery.add_peer("http://refresh.me", source="test")
+    peer = discovery.get_discovered_peers()[0]
+    first_seen = peer.last_seen
+
+    # Wait a moment to ensure a measurable delta
+    await asyncio.sleep(0.01)
+
+    refreshed = await discovery.refresh_peer("http://refresh.me")
+    assert refreshed is True
+    assert peer.last_seen > first_seen

tests/unit/mcpgateway/test_forward.py

@@ -0,0 +1,172 @@
+# -*- coding: utf-8 -*-
+"""Unit tests for Federation Forwarding Service.
+
+Copyright 2025
+SPDX-License-Identifier: Apache-2.0
+Authors: Mihai Criveti
+
+Comprehensive unit tests for the forwarding service module.
+"""
+
+import asyncio
+from datetime import datetime
+
+import pytest
+
+from mcpgateway.federation.forward import ForwardingError, ForwardingService, ToolResult
+from mcpgateway.types import TextContent
+
+# ---------------------------------------------------------------------------
+# Tiny dummy ORM objects + fake Session
+# ---------------------------------------------------------------------------
+
+
+class DummyGateway:
+    def __init__(self, id_, name, url, is_active=True):
+        self.id = id_
+        self.name = name
+        self.url = url
+        self.is_active = is_active
+        self.last_seen: datetime | None = None
+
+
+class DummyTool:
+    def __init__(self, id_, name, gateway_id=None, is_active=True):
+        self.id = id_
+        self.name = name
+        self.gateway_id = gateway_id
+        self.is_active = is_active
+
+
+class _FakeResult:
+    def __init__(self, scalar=None, scalar_list=None):
+        self._scalar = scalar
+        self._scalar_list = scalar_list or []
+
+    def scalar_one_or_none(self):
+        return self._scalar
+
+    class _Proxy:
+        def __init__(self, items):
+            self._items = items
+
+        def all(self):
+            return self._items
+
+    def scalars(self):
+        return self._Proxy(self._scalar_list)
+
+
+class FakeSession:
+    def __init__(self, gateways=None, tools=None):
+        self._gateways = gateways or []
+        self._tools = tools or []
+
+    def get(self, _model, pk):
+        for gw in self._gateways:
+            if gw.id == pk:
+                return gw
+        return None
+
+    def execute(self, query):  # pragma: no cover
+        txt = str(query).lower()
+        if "tool" in txt:
+            return _FakeResult(scalar=self._tools[0] if self._tools else None)
+        return _FakeResult(scalar_list=self._gateways)
+
+
+# ---------------------------------------------------------------------------
+# ForwardingService fixture with network stubbed out
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+async def fwd_service(monkeypatch):
+    svc = ForwardingService()
+
+    class _FakeResp:
+        def __init__(self, payload):
+            self._payload = payload
+
+        def raise_for_status(self):
+            pass
+
+        def json(self):
+            return self._payload
+
+    async def fake_post(url, json=None, headers=None):  # noqa: D401
+        return _FakeResp({"result": {"method": json["method"]}})
+
+    monkeypatch.setattr(svc._http_client, "post", fake_post)
+
+    from mcpgateway.config import settings
+
+    monkeypatch.setattr(settings, "max_tool_retries", 1, raising=False)
+    yield svc
+    await svc.stop()
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.anyio
+async def test_forward_to_gateway_success(fwd_service):
+    gw = DummyGateway(1, "Alpha", "http://alpha")
+    db = FakeSession(gateways=[gw])
+
+    result = await fwd_service._forward_to_gateway(db, 1, "ping", {"x": 1})
+    assert result == {"method": "ping"}
+    assert isinstance(gw.last_seen, datetime)
+
+
+@pytest.mark.anyio
+async def test_forward_tool_request_parses_result(monkeypatch, fwd_service):
+    # Fake gateway call to produce valid ContentType payload
+    async def fake_forward(db, gid, method, params):  # noqa: D401
+        assert method == "tools/invoke"
+        payload = {
+            "content": [TextContent(type="text", text="OK").model_dump()],
+            "is_error": False,
+        }
+        return payload
+
+    monkeypatch.setattr(fwd_service, "_forward_to_gateway", fake_forward)  # type: ignore[arg-type]
+
+    tool = DummyTool(1, "echo", gateway_id=42)
+    db = FakeSession(gateways=[DummyGateway(42, "EchoGW", "http://echo")], tools=[tool])
+
+    result: ToolResult = await fwd_service.forward_tool_request(db, "echo", {"msg": "hi"})
+    assert result.is_error is False
+    assert isinstance(result.content[0], TextContent) and result.content[0].text == "OK"
+
+
+@pytest.mark.anyio
+async def test_rate_limit(monkeypatch):
+    from mcpgateway.config import settings
+
+    monkeypatch.setattr(settings, "tool_rate_limit", 2, raising=False)
+
+    svc = ForwardingService()
+    url = "http://beta"
+    assert svc._check_rate_limit(url)
+    assert svc._check_rate_limit(url)
+    assert svc._check_rate_limit(url) is False  # third call exceeds limit
+
+
+@pytest.mark.anyio
+async def test_forward_to_all_partial_success(monkeypatch, fwd_service):
+    gw_ok = DummyGateway(1, "GoodGW", "http://good")
+    gw_bad = DummyGateway(2, "BadGW", "http://bad")
+
+    async def fake_forward(db, gid, method, params):  # noqa: D401
+        if gid == 1:
+            return "ok!"
+        raise ForwardingError("boom")
+
+    monkeypatch.setattr(fwd_service, "_forward_to_gateway", fake_forward)  # type: ignore[arg-type]
+
+    db = FakeSession(gateways=[gw_ok, gw_bad])
+    results = await fwd_service._forward_to_all(db, "stats/get")
+    assert results == ["ok!"]