test case update (#224)

* test case update

* Test case update

Signed-off-by: Mohan Lakshmaiah <[email protected]>

---------

Signed-off-by: Mohan Lakshmaiah <[email protected]>
Co-authored-by: Mohan Lakshmaiah <[email protected]>

Author: MohanLaksh

tests/unit/mcpgateway/handlers/test_sampling.py

@@ -14,6 +14,7 @@
 
 No external MCP server is started; we test the isolated utility pieces that
 have no heavy dependencies.
+
 """
 
 # Future
@@ -61,6 +62,70 @@ async def collector(msg):
     assert len(sent) == 1 and sent[0].message["id"] == 2
     assert sent[0].event_id == eid2
 
+@pytest.mark.asyncio
+async def test_event_store_no_new_events():
+    store = InMemoryEventStore(max_events_per_stream=10)
+    stream_id = "stream1"
+    eid = await store.store_event(stream_id, {"val": 42})
+    sent = []
+    async def collector(msg):
+        sent.append(msg)
+
+    returned = await store.replay_events_after(eid, collector)
+    assert returned == stream_id
+    # No new events were stored, so nothing should be sent
+    assert sent == []
+
+@pytest.mark.asyncio
+async def test_event_store_multiple_replay():
+    store = InMemoryEventStore(max_events_per_stream=10)
+    stream_id = "stream1"
+    # Store three events
+    eids = []
+    for i in range(3):
+        eids.append(await store.store_event(stream_id, {"n": i}))
+    sent = []
+    async def collector(msg):
+        sent.append(msg)
+
+    # Replay after the first event should get the 2nd and 3rd
+    returned = await store.replay_events_after(eids[0], collector)
+    assert returned == stream_id
+    assert [msg.message["n"] for msg in sent] == [1, 2]
+
+@pytest.mark.asyncio
+async def test_event_store_cross_streams():
+    store = InMemoryEventStore(max_events_per_stream=10)
+    s1, s2 = "s1", "s2"
+    # Store events in two different streams
+    eid1_s1 = await store.store_event(s1, {"val": 1})
+    eid1_s2 = await store.store_event(s2, {"val": 2})
+    eid2_s1 = await store.store_event(s1, {"val": 3})
+    sent = []
+    async def collector(msg):
+        sent.append(msg)
+
+    # Replay on stream s1 after its first ID
+    returned = await store.replay_events_after(eid1_s1, collector)
+    assert returned == s1
+    # Should only get the event from s1 (val=3), not the s2 event
+    assert [msg.message["val"] for msg in sent] == [3]
+
+@pytest.mark.asyncio
+async def test_event_store_eviction_of_oldest():
+    store = InMemoryEventStore(max_events_per_stream=1)
+    stream_id = "s"
+    eid_old = await store.store_event(stream_id, {"x": "old"})
+    # Storing a second event evicts the first (due to maxlen=1):contentReference[oaicite:8]{index=8}:contentReference[oaicite:9]{index=9}
+    await store.store_event(stream_id, {"x": "new"})
+    sent = []
+    async def collector(msg):
+        sent.append(msg)
+
+    result = await store.replay_events_after(eid_old, collector)
+    # The first event ID has been evicted, so it should not be found
+    assert result is None
+    assert sent == []
 
 @pytest.mark.asyncio
 async def test_event_store_eviction():
@@ -147,6 +212,42 @@ async def send(msg):
     assert sent and sent[0]["type"] == "http.response.start"
     assert sent[0]["status"] == tr.HTTP_401_UNAUTHORIZED
 
+@pytest.mark.asyncio
+async def test_auth_valid_token(monkeypatch):
+    # Simulate verify_credentials always succeeding
+    async def fake_verify(token):
+        assert token == "good-token"
+        return {"ok": True}
+    monkeypatch.setattr(tr, "verify_credentials", fake_verify)
+
+    messages = []
+    async def send(msg):
+        messages.append(msg)
+
+    scope = _make_scope("/servers/1/mcp",
+                        headers=[(b"authorization", b"Bearer good-token")])
+    assert await streamable_http_auth(scope, None, send) is True
+    assert messages == []  # No response sent on success
+
+@pytest.mark.asyncio
+async def test_auth_invalid_token_raises(monkeypatch):
+    # Simulate verify_credentials raising (invalid token scenario)
+    async def fake_verify(token):
+        raise ValueError("bad token")
+    monkeypatch.setattr(tr, "verify_credentials", fake_verify)
+
+    sent = []
+    async def send(msg):
+        sent.append(msg)
+
+    scope = _make_scope("/servers/1/mcp",
+                        headers=[(b"authorization", b"Bearer bad-token")])
+    result = await streamable_http_auth(scope, None, send)
+    assert result is False
+    # Expect an HTTP 401 response to be sent
+    assert sent and sent[0]["type"] == "http.response.start"
+    assert sent[0]["status"] == tr.HTTP_401_UNAUTHORIZED
+
 
 # ---------------------------------------------------------------------------
 # SamplingHandler tests
@@ -185,6 +286,13 @@ async def test_select_model_by_hint(handler):
 
     assert handler._select_model(prefs) == "claude-3-sonnet"  # pylint: disable=protected-access
 
+@pytest.mark.asyncio
+async def test_select_model_no_suitable_model(handler):
+    # Remove all supported models to force error
+    handler._supported_models = {}
+    prefs = _t.SimpleNamespace(hints=[], cost_priority=0, speed_priority=0, intelligence_priority=0)
+    with pytest.raises(SamplingError):
+        handler._select_model(prefs)
 
 # ---------------------------------------------------------------------------
 # _validate_message
@@ -203,6 +311,54 @@ def test_validate_message(handler):
     assert handler._validate_message(valid_image)  # pylint: disable=protected-access
     assert not handler._validate_message(invalid)  # pylint: disable=protected-access
 
+def test_validate_message_missing_image_fields(handler):
+    # Missing 'data' field in image content
+    invalid_img1 = {"role": "assistant", "content": {"type": "image", "mime_type": "image/png"}}
+    # Missing 'mime_type' field
+    invalid_img2 = {"role": "assistant", "content": {"type": "image", "data": "AAA"}}
+    # Unknown content type
+    invalid_img3 = {"role": "user", "content": {"type": "audio", "data": "xxx"}}
+
+    assert not handler._validate_message(invalid_img1)
+    assert not handler._validate_message(invalid_img2)
+    assert not handler._validate_message(invalid_img3)
+
+@pytest.mark.asyncio
+async def test_add_context_returns_messages(handler):
+    # Should just return the messages as-is (stub)
+    msgs = [{"role": "user", "content": {"type": "text", "text": "hi"}}]
+    result = await handler._add_context(None, msgs, "irrelevant")
+    assert result == msgs
+
+def test_mock_sample_no_user_message(handler):
+    # No user message in the list
+    msgs = [{"role": "assistant", "content": {"type": "text", "text": "hi"}}]
+    result = handler._mock_sample(msgs)
+    assert "I'm not sure" in result
+
+def test_mock_sample_image_message(handler):
+    # Last user message is image
+    msgs = [
+        {"role": "user", "content": {"type": "image", "data": "xxx", "mime_type": "image/png"}}
+    ]
+    result = handler._mock_sample(msgs)
+    assert "I see the image" in result
+
+def test_validate_message_invalid_role(handler):
+    msg = {"role": "system", "content": {"type": "text", "text": "hi"}}
+    assert not handler._validate_message(msg)
+
+def test_validate_message_missing_content(handler):
+    msg = {"role": "user"}
+    assert not handler._validate_message(msg)
+
+def test_validate_message_exception_path(handler):
+    # Simulate exception in validation
+    class BadDict(dict):
+        def get(self, k, d=None):
+            raise Exception("fail")
+    msg = {"role": "user", "content": BadDict()}
+    assert not handler._validate_message(msg)
 
 # ---------------------------------------------------------------------------
 # create_message success + error paths
@@ -228,6 +384,29 @@ async def test_create_message_success(monkeypatch, handler):
     assert result.role == sp.Role.ASSISTANT
     assert result.content.text.startswith("You said: Hello")
 
+@pytest.mark.asyncio
+async def test_create_message_multiple_user_messages(monkeypatch, handler):
+    # Return neutral preferences with no hints
+    neutral_prefs = _t.SimpleNamespace(
+        hints=[], cost_priority=0.5, speed_priority=0.3, intelligence_priority=0.2
+    )
+    monkeypatch.setattr(sp.ModelPreferences, "model_validate", lambda x: neutral_prefs)
+
+    # Conversation with an assistant message and then a user message
+    request = {
+        "messages": [
+            {"role": "assistant", "content": {"type": "text", "text": "Hi"}},
+            {"role": "user", "content": {"type": "text", "text": "Hello"}}
+        ],
+        "maxTokens": 10,
+        "modelPreferences": {}
+    }
+
+    result = await handler.create_message(db=None, request=request)
+    assert result.role == sp.Role.ASSISTANT
+    # The response should reference the last user message "Hello"
+    assert "You said: Hello" in result.content.text
+
 
 @pytest.mark.asyncio
 async def test_create_message_no_messages(monkeypatch, handler):
@@ -237,3 +416,45 @@ async def test_create_message_no_messages(monkeypatch, handler):
 
     with pytest.raises(SamplingError):
         await handler.create_message(db=None, request=request)
+
+@pytest.mark.asyncio
+async def test_create_message_raises_on_no_user_message(monkeypatch, handler):
+    # Even if there are assistant messages, at least one user message is required
+    monkeypatch.setattr(sp.ModelPreferences, "model_validate",
+                        lambda x: _t.SimpleNamespace(hints=[], cost_priority=0, speed_priority=0, intelligence_priority=0))
+    request = {"messages": [], "maxTokens": 5, "modelPreferences": {}}
+    with pytest.raises(SamplingError):
+        await handler.create_message(db=None, request=request)
+
+@pytest.mark.asyncio
+async def test_create_message_missing_max_tokens(monkeypatch, handler):
+    monkeypatch.setattr(sp.ModelPreferences, "model_validate", lambda _x: _t.SimpleNamespace(hints=[], cost_priority=0, speed_priority=0, intelligence_priority=0))
+    request = {"messages": [{"role": "user", "content": {"type": "text", "text": "hi"}}]}
+    with pytest.raises(SamplingError):
+        await handler.create_message(db=None, request=request)
+
+@pytest.mark.asyncio
+async def test_create_message_invalid_message(monkeypatch, handler):
+    monkeypatch.setattr(sp.ModelPreferences, "model_validate", lambda _x: _t.SimpleNamespace(hints=[], cost_priority=0, speed_priority=0, intelligence_priority=0))
+    # Invalid message: missing text
+    request = {
+        "messages": [{"role": "user", "content": {"type": "text"}}],
+        "maxTokens": 5,
+        "modelPreferences": {},
+    }
+    with pytest.raises(SamplingError):
+        await handler.create_message(db=None, request=request)
+
+@pytest.mark.asyncio
+async def test_create_message_exception_propagation(monkeypatch, handler):
+    # Patch _select_model to raise
+    monkeypatch.setattr(handler, "_select_model", lambda prefs: (_ for _ in ()).throw(Exception("fail")))
+    monkeypatch.setattr(sp.ModelPreferences, "model_validate", lambda _x: _t.SimpleNamespace(hints=[], cost_priority=0, speed_priority=0, intelligence_priority=0))
+    request = {
+        "messages": [{"role": "user", "content": {"type": "text", "text": "hi"}}],
+        "maxTokens": 5,
+        "modelPreferences": {},
+    }
+    with pytest.raises(SamplingError) as exc:
+        await handler.create_message(db=None, request=request)
+    assert "fail" in str(exc.value)

tests/unit/mcpgateway/transports/test_sse_transport.py

@@ -6,12 +6,14 @@
 Authors: Mihai Criveti
 
 Tests for the MCP Gateway SSE transport implementation.
+
 """
 
 # Standard
 import asyncio
 import json
-from unittest.mock import Mock
+import types
+from unittest.mock import Mock, AsyncMock, patch
 
 # First-Party
 from mcpgateway.transports.sse_transport import SSETransport
@@ -81,6 +83,51 @@ async def test_send_message_not_connected(self, sse_transport):
         # Should raise error
         with pytest.raises(RuntimeError, match="Transport not connected"):
             await sse_transport.send_message(message)
+    @pytest.mark.asyncio
+    async def test_receive_message_not_connected(self, sse_transport):
+        """receive_message should raise RuntimeError if not connected."""
+        with pytest.raises(RuntimeError):
+            async for _ in sse_transport.receive_message():
+                pass
+
+    @pytest.mark.asyncio
+    async def test_send_message_queue_exception(self, sse_transport):
+        """send_message should log and raise if queue.put fails."""
+        await sse_transport.connect()
+        with patch.object(sse_transport._message_queue, "put", side_effect=Exception("fail")), \
+            patch("mcpgateway.transports.sse_transport.logger") as mock_logger:
+            with pytest.raises(Exception, match="fail"):
+                await sse_transport.send_message({"foo": "bar"})
+            assert mock_logger.error.called
+
+    @pytest.mark.asyncio
+    async def test_receive_message_cancelled(self, sse_transport):
+        """Test receive_message handles CancelledError and logs."""
+        await sse_transport.connect()
+        with patch("asyncio.sleep", side_effect=asyncio.CancelledError), \
+            patch("mcpgateway.transports.sse_transport.logger") as mock_logger:
+            gen = sse_transport.receive_message()
+            await gen.__anext__()  # initialize message
+            with pytest.raises(asyncio.CancelledError):
+                await gen.__anext__()
+            # Check that logger.info was called with the cancel message
+            assert any(
+                "SSE receive loop cancelled" in str(call)
+                for call in [args[0] for args, _ in mock_logger.info.call_args_list]
+            )
+    @pytest.mark.asyncio
+    async def test_receive_message_finally_logs(self, sse_transport):
+        """Test receive_message logs in finally block."""
+        await sse_transport.connect()
+        with patch("asyncio.sleep", side_effect=Exception("fail")), \
+            patch("mcpgateway.transports.sse_transport.logger") as mock_logger:
+            gen = sse_transport.receive_message()
+            await gen.__anext__()  # initialize message
+            with pytest.raises(Exception):
+                await gen.__anext__()
+            assert any("SSE receive loop ended" in str(call) for call in mock_logger.info.call_args_list)
+
+
 
     @pytest.mark.asyncio
     async def test_create_sse_response(self, sse_transport, mock_request):
@@ -99,6 +146,37 @@ async def test_create_sse_response(self, sse_transport, mock_request):
         assert response.headers["Cache-Control"] == "no-cache"
         assert response.headers["Content-Type"] == "text/event-stream"
         assert response.headers["X-MCP-SSE"] == "true"
+    
+    @pytest.mark.asyncio
+    async def test_create_sse_response_event_generator_error(self, sse_transport, mock_request):
+        """Test event_generator handles generic Exception and CancelledError."""
+        await sse_transport.connect()
+        # Patch _message_queue.get to raise Exception, then CancelledError
+        with patch.object(sse_transport._message_queue, "get", side_effect=[Exception("fail"), asyncio.CancelledError()]), \
+             patch("mcpgateway.transports.sse_transport.logger") as mock_logger:
+            response = await sse_transport.create_sse_response(mock_request)
+            gen = response.body_iterator
+            await gen.__anext__()  # endpoint
+            await gen.__anext__()  # keepalive
+            # Should yield error event
+            event = await gen.__anext__()
+            assert event["event"] == "error"
+            assert "fail" in event["data"]
+            # Should handle CancelledError gracefully and stop
+            with pytest.raises(StopAsyncIteration):
+                await gen.__anext__()
+            assert mock_logger.error.called or mock_logger.info.called
+            
+    def test_session_id_property(self, sse_transport):
+        """Test session_id property returns the correct value."""
+        assert sse_transport.session_id == sse_transport._session_id
+
+    @pytest.mark.asyncio
+    async def test_client_disconnected(self, sse_transport, mock_request):
+        """Test _client_disconnected returns correct state."""
+        assert await sse_transport._client_disconnected(mock_request) is False
+        sse_transport._client_gone.set()
+        assert await sse_transport._client_disconnected(mock_request) is True
 
     @pytest.mark.asyncio
     async def test_receive_message(self, sse_transport):