Merge pull request #1459 from MohanLaksh/test_cases

Test cases

Author: madhav165

tests/unit/mcpgateway/instrumentation/test_sqlalchemy.py

@@ -0,0 +1,179 @@
+# -*- coding: utf-8 -*-
+"""Location: ./tests/unit/mcpgateway/instrumentation/test_sqlalchemy.py
+Copyright 2025
+SPDX-License-Identifier: Apache-2.0
+Authors: Mihai Criveti
+
+Unit tests for sqlalchemy instrumentation.
+"""
+import pytest
+import threading
+import queue
+import time
+from unittest.mock import MagicMock, patch
+
+import mcpgateway.instrumentation.sqlalchemy as sa
+
+
+@pytest.fixture(autouse=True)
+def reset_globals():
+    sa._query_tracking.clear()
+    sa._instrumentation_context = threading.local()
+    sa._span_queue = queue.Queue(maxsize=2)
+    sa._shutdown_event.clear()
+    sa._span_writer_thread = None
+    yield
+    sa._query_tracking.clear()
+
+
+def test_before_cursor_execute_stores_tracking():
+    conn = MagicMock()
+    sa._before_cursor_execute(conn, None, "SELECT * FROM test", {"id": 1}, None, False)
+    conn_id = id(conn)
+    assert conn_id in sa._query_tracking
+    tracking = sa._query_tracking[conn_id]
+    assert tracking["statement"] == "SELECT * FROM test"
+    assert "start_time" in tracking
+
+
+def test_after_cursor_execute_no_tracking():
+    conn = MagicMock()
+    sa._after_cursor_execute(conn, None, "SELECT * FROM test", None, None, False)
+    # Should not raise or enqueue anything
+    assert sa._span_queue.empty()
+
+
+def test_after_cursor_execute_inside_span_creation_skips():
+    conn = MagicMock()
+    conn_id = id(conn)
+    sa._query_tracking[conn_id] = {"start_time": time.time(), "statement": "SELECT 1", "parameters": None, "executemany": False}
+    sa._instrumentation_context.inside_span_creation = True
+    sa._after_cursor_execute(conn, MagicMock(), "SELECT 1", None, None, False)
+    assert sa._span_queue.empty()
+
+
+def test_after_cursor_execute_observability_table_skips(caplog):
+    import logging
+    caplog.set_level(logging.DEBUG)
+    sa.logger.setLevel(logging.DEBUG)
+    sa.logger.propagate = True
+    conn = MagicMock()
+    conn_id = id(conn)
+    sa._query_tracking[conn_id] = {"start_time": time.time(), "statement": "SELECT * FROM observability_spans", "parameters": None, "executemany": False}
+    sa._after_cursor_execute(conn, MagicMock(), "SELECT * FROM observability_spans", None, None, False)
+    assert "Skipping instrumentation" in caplog.text
+
+
+def test_after_cursor_execute_with_trace_id_calls_create_query_span():
+    conn = MagicMock()
+    conn.info = {"trace_id": "abc123"}
+    conn_id = id(conn)
+    sa._query_tracking[conn_id] = {"start_time": time.time(), "statement": "SELECT * FROM users", "parameters": None, "executemany": False}
+    with patch.object(sa, "_create_query_span") as mock_create:
+        sa._after_cursor_execute(conn, MagicMock(rowcount=5), "SELECT * FROM users", None, None, False)
+        mock_create.assert_called_once()
+        args, kwargs = mock_create.call_args
+        assert kwargs["trace_id"] == "abc123"
+
+
+def test_after_cursor_execute_without_trace_id_logs_debug(caplog):
+    import logging
+    caplog.set_level(logging.DEBUG)
+    sa.logger.setLevel(logging.DEBUG)
+    sa.logger.propagate = True
+    conn = MagicMock()
+    conn.info = {}
+    conn_id = id(conn)
+    sa._query_tracking[conn_id] = {"start_time": time.time(), "statement": "SELECT * FROM users", "parameters": None, "executemany": False}
+    sa._after_cursor_execute(conn, MagicMock(rowcount=5), "SELECT * FROM users", None, None, False)
+    assert "Query executed without trace context" in caplog.text
+
+
+def test_create_query_span_enqueues_successfully(caplog):
+    import logging
+    caplog.set_level(logging.DEBUG)
+    sa.logger.setLevel(logging.DEBUG)
+    sa.logger.propagate = True
+    sa._create_query_span("trace123", "SELECT * FROM test", 10.0, 5, False)
+    assert not sa._span_queue.empty()
+    assert "Enqueued span" in caplog.text
+
+
+def test_create_query_span_queue_full_warns(caplog):
+    sa._span_queue = queue.Queue(maxsize=1)
+    sa._span_queue.put({"dummy": "data"})
+    sa._create_query_span("trace123", "SELECT * FROM test", 10.0, 5, False)
+    assert "Span queue is full" in caplog.text
+
+
+def test_create_query_span_exception_does_not_raise(caplog):
+    import logging
+    caplog.set_level(logging.DEBUG)
+    sa.logger.setLevel(logging.DEBUG)
+    sa.logger.propagate = True
+    with patch("mcpgateway.instrumentation.sqlalchemy._span_queue.put_nowait", side_effect=Exception("fail")):
+        sa._create_query_span("trace123", "SELECT * FROM test", 10.0, 5, False)
+    assert "Failed to enqueue query span" in caplog.text
+
+
+def test_write_span_to_db_success():
+    span_data = {
+        "trace_id": "t1",
+        "name": "db.query.select",
+        "kind": "client",
+        "resource_type": "database",
+        "resource_name": "SELECT",
+        "start_attributes": {},
+        "end_attributes": {},
+        "status": "ok",
+        "duration_ms": 10.0,
+        "row_count": 1,
+    }
+    mock_service = MagicMock()
+    mock_db = MagicMock()
+    mock_span = MagicMock()
+    mock_db.query().filter_by().first.return_value = mock_span
+    with patch("mcpgateway.services.observability_service.ObservabilityService", return_value=mock_service), \
+         patch("mcpgateway.db.SessionLocal", return_value=mock_db), \
+         patch("mcpgateway.db.ObservabilitySpan", MagicMock()):
+        sa._write_span_to_db(span_data)
+    mock_service.start_span.assert_called_once()
+    mock_service.end_span.assert_called_once()
+    mock_db.commit.assert_called_once()
+
+
+def test_write_span_to_db_exception_logs_warning(caplog):
+    with patch("mcpgateway.services.observability_service.ObservabilityService", side_effect=Exception("fail")):
+        sa._write_span_to_db({})
+    assert "Failed to write query span" in caplog.text
+
+
+def test_span_writer_worker_processes_queue(monkeypatch):
+    span_data = {"trace_id": "t1", "name": "db.query.select", "kind": "client", "resource_type": "database", "resource_name": "SELECT", "start_attributes": {}, "end_attributes": {}, "status": "ok", "duration_ms": 10.0}
+    sa._span_queue.put(span_data)
+    mock_write = MagicMock()
+    monkeypatch.setattr(sa, "_write_span_to_db", mock_write)
+    thread = threading.Thread(target=lambda: (time.sleep(0.1), sa._shutdown_event.set()))
+    thread.start()
+    sa._span_writer_worker()
+    mock_write.assert_called_once()
+
+
+def test_instrument_sqlalchemy_starts_thread_and_registers_events():
+    engine = MagicMock()
+    with patch("mcpgateway.instrumentation.sqlalchemy.event.listen") as mock_listen, \
+         patch("mcpgateway.instrumentation.sqlalchemy.threading.Thread") as mock_thread:
+        mock_thread.return_value = MagicMock(is_alive=lambda: False)
+        sa.instrument_sqlalchemy(engine)
+        assert mock_listen.call_count == 2
+        mock_thread.assert_called_once()
+
+
+def test_attach_trace_to_session_sets_trace_id():
+    session = MagicMock()
+    connection = MagicMock()
+    connection.info = {}
+    session.bind = True
+    session.connection.return_value = connection
+    sa.attach_trace_to_session(session, "trace123")
+    assert connection.info["trace_id"] == "trace123"

tests/unit/mcpgateway/middleware/test_auth_middleware.py

@@ -0,0 +1,146 @@
+# -*- coding: utf-8 -*-
+"""Location: ./tests/unit/mcpgateway/middleware/test_auth_middleware.py
+Copyright 2025
+SPDX-License-Identifier: Apache-2.0
+Authors: Mihai Criveti
+
+Unit tests for auth middleware.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+from starlette.requests import Request
+from starlette.responses import Response
+from mcpgateway.middleware.auth_middleware import AuthContextMiddleware
+
+
+@pytest.mark.asyncio
+async def test_health_and_static_paths_skipped(monkeypatch):
+    """Ensure middleware skips health and static paths."""
+    middleware = AuthContextMiddleware(app=AsyncMock())
+    call_next = AsyncMock(return_value=Response("ok"))
+
+    for path in ["/health", "/healthz", "/ready", "/metrics", "/static/logo.png"]:
+        request = MagicMock(spec=Request)
+        request.url.path = path
+        response = await middleware.dispatch(request, call_next)
+        call_next.assert_awaited_once_with(request)
+        assert response.status_code == 200
+        call_next.reset_mock()
+
+
+@pytest.mark.asyncio
+async def test_no_token_continues(monkeypatch):
+    """If no token found, request continues without user context."""
+    middleware = AuthContextMiddleware(app=AsyncMock())
+    call_next = AsyncMock(return_value=Response("ok"))
+    request = MagicMock(spec=Request)
+    request.url.path = "/api/data"
+    request.cookies = {}
+    request.headers = {}
+
+    response = await middleware.dispatch(request, call_next)
+    call_next.assert_awaited_once_with(request)
+    assert response.status_code == 200
+    # request.state is a MagicMock, so user may exist as mock attribute
+    # Instead, ensure user was never set explicitly
+    # Ensure user attribute was not explicitly set (MagicMock defaults to having attributes)
+    assert "user" not in request.state.__dict__
+
+
+@pytest.mark.asyncio
+async def test_token_from_cookie(monkeypatch):
+    """Token extracted from cookie triggers authentication."""
+    middleware = AuthContextMiddleware(app=AsyncMock())
+    call_next = AsyncMock(return_value=Response("ok"))
+    request = MagicMock(spec=Request)
+    request.url.path = "/api/data"
+    request.cookies = {"jwt_token": "cookie_token"}
+    request.headers = {}
+
+    mock_user = MagicMock()
+    mock_user.email = "user@example.com"
+
+    with patch("mcpgateway.middleware.auth_middleware.SessionLocal", return_value=MagicMock()) as mock_session, \
+         patch("mcpgateway.middleware.auth_middleware.get_current_user", AsyncMock(return_value=mock_user)):
+        response = await middleware.dispatch(request, call_next)
+
+    call_next.assert_awaited_once_with(request)
+    assert response.status_code == 200
+    assert request.state.user.email == "user@example.com"
+    mock_session.return_value.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_token_from_header(monkeypatch):
+    """Token extracted from Authorization header triggers authentication."""
+    middleware = AuthContextMiddleware(app=AsyncMock())
+    call_next = AsyncMock(return_value=Response("ok"))
+    request = MagicMock(spec=Request)
+    request.url.path = "/api/data"
+    request.cookies = {}
+    request.headers = {"authorization": "Bearer header_token"}
+
+    mock_user = MagicMock()
+    mock_user.email = "header@example.com"
+
+    with patch("mcpgateway.middleware.auth_middleware.SessionLocal", return_value=MagicMock()) as mock_session, \
+         patch("mcpgateway.middleware.auth_middleware.get_current_user", AsyncMock(return_value=mock_user)):
+        response = await middleware.dispatch(request, call_next)
+
+    call_next.assert_awaited_once_with(request)
+    assert response.status_code == 200
+    assert request.state.user.email == "header@example.com"
+    mock_session.return_value.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_authentication_failure(monkeypatch):
+    """Authentication failure should log and continue without user context."""
+    middleware = AuthContextMiddleware(app=AsyncMock())
+    call_next = AsyncMock(return_value=Response("ok"))
+    request = MagicMock(spec=Request)
+    request.url.path = "/api/data"
+    request.cookies = {"jwt_token": "bad_token"}
+    request.headers = {}
+
+    with patch("mcpgateway.middleware.auth_middleware.SessionLocal", return_value=MagicMock()) as mock_session, \
+         patch("mcpgateway.middleware.auth_middleware.get_current_user", AsyncMock(side_effect=Exception("Invalid token"))), \
+         patch("mcpgateway.middleware.auth_middleware.logger") as mock_logger:
+        response = await middleware.dispatch(request, call_next)
+
+    call_next.assert_awaited_once_with(request)
+    assert response.status_code == 200
+    # Ensure user attribute was not explicitly set (MagicMock defaults to having attributes)
+    assert "user" not in request.state.__dict__
+    # Verify log message contains failure text
+    logged_messages = [args[0] for args, _ in mock_logger.info.call_args_list]
+    assert any("✗ Auth context extraction failed" in msg for msg in logged_messages)
+    mock_session.return_value.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_db_close_exception(monkeypatch):
+    """Ensure db.close exceptions are logged but do not break flow."""
+    middleware = AuthContextMiddleware(app=AsyncMock())
+    call_next = AsyncMock(return_value=Response("ok"))
+    request = MagicMock(spec=Request)
+    request.url.path = "/api/data"
+    request.cookies = {"jwt_token": "token"}
+    request.headers = {}
+
+    mock_user = MagicMock()
+    mock_user.email = "user@example.com"
+    mock_db = MagicMock()
+    mock_db.close.side_effect = Exception("close error")
+
+    with patch("mcpgateway.middleware.auth_middleware.SessionLocal", return_value=mock_db), \
+         patch("mcpgateway.middleware.auth_middleware.get_current_user", AsyncMock(return_value=mock_user)), \
+         patch("mcpgateway.middleware.auth_middleware.logger") as mock_logger:
+        response = await middleware.dispatch(request, call_next)
+
+    call_next.assert_awaited_once_with(request)
+    assert response.status_code == 200
+    # Verify log message contains close error text
+    logged_debugs = [args[0] for args, _ in mock_logger.debug.call_args_list]
+    assert any("Failed to close database session" in msg for msg in logged_debugs)