test: fix and improve unit tests with comprehensive coverage

- Fix API key mocking using proper Settings._get_api_key_from_args patch
- Fix async mock issues by using Mock for sync methods (json, raise_for_status)
- Add comprehensive tests for pia_search_facets tool
- Fix server exception test with proper HTTP client mocking
- Improve test coverage to 71% overall, 97% for pia_search, 89% for facets
- Test actual functionality: HTTP integration, JSON parsing, error handling
- Ensure proper MCP protocol compliance testing
- Add tests for API key validation, HTTP errors, and API error responses

All 14 tests now pass with proper mocking of external dependencies

Author: dividor

tests/test_server.py

@@ -37,10 +37,16 @@ async def test_call_unknown_tool():
 async def test_call_tool_exception():
     """Test that exceptions in tool calls are handled properly."""
     with patch(
-        "pia_mcp_server.tools.handle_pia_search", side_effect=Exception("Test error")
+        "pia_mcp_server.config.Settings._get_api_key_from_args", return_value="test_key"
     ):
-        result = await call_tool("pia_search", {"query": "test"})
+        with patch("httpx.AsyncClient") as mock_client:
+            # Make the client throw an exception when post is called
+            mock_client_instance = AsyncMock()
+            mock_client_instance.post.side_effect = Exception("Test error")
+            mock_client.return_value.__aenter__.return_value = mock_client_instance
 
-        assert len(result) == 1
-        assert result[0].type == "text"
-        assert "Error: Test error" in result[0].text
+            result = await call_tool("pia_search", {"query": "test"})
+
+            assert len(result) == 1
+            assert result[0].type == "text"
+            assert "Error: Test error" in result[0].text

tests/test_tools.py

@@ -1,9 +1,10 @@
 """Tests for tools module."""
 
 import pytest
-from unittest.mock import AsyncMock, patch, PropertyMock
+from unittest.mock import AsyncMock, patch, PropertyMock, Mock
 import httpx
 from pia_mcp_server.tools.pia_search import handle_pia_search
+from pia_mcp_server.tools.pia_search_facets import handle_pia_search_facets
 from pia_mcp_server.config import Settings
 
 settings = Settings()
@@ -12,10 +13,7 @@
 @pytest.mark.asyncio
 async def test_pia_search_no_api_key():
     """Test PIA search without API key."""
-    with patch(
-        "pia_mcp_server.tools.pia_search.settings.API_KEY",
-        side_effect=ValueError("PIA API key is required"),
-    ):
+    with patch.object(Settings, "_get_api_key_from_args", return_value=None):
         result = await handle_pia_search({"query": "test"})
 
         assert len(result) == 1
@@ -36,16 +34,15 @@ async def test_pia_search_success():
         },
     }
 
-    with patch.object(
-        type(settings),
-        "API_KEY",
-        new_callable=lambda: PropertyMock(return_value="test_key"),
-    ):
-        with patch("httpx.AsyncClient.post") as mock_post:
-            mock_response_obj = AsyncMock()
+    with patch.object(Settings, "_get_api_key_from_args", return_value="test_key"):
+        with patch("httpx.AsyncClient") as mock_client:
+            mock_response_obj = Mock()
             mock_response_obj.json.return_value = mock_response
-            mock_response_obj.raise_for_status = AsyncMock()
-            mock_post.return_value = mock_response_obj
+            mock_response_obj.raise_for_status.return_value = None
+
+            mock_client_instance = AsyncMock()
+            mock_client_instance.post.return_value = mock_response_obj
+            mock_client.return_value.__aenter__.return_value = mock_client_instance
 
             result = await handle_pia_search({"query": "test fraud"})
 
@@ -62,16 +59,15 @@ async def test_pia_search_api_error():
         "error": {"code": -32000, "message": "Invalid API key"},
     }
 
-    with patch.object(
-        type(settings),
-        "API_KEY",
-        new_callable=lambda: PropertyMock(return_value="invalid_key"),
-    ):
-        with patch("httpx.AsyncClient.post") as mock_post:
-            mock_response_obj = AsyncMock()
+    with patch.object(Settings, "_get_api_key_from_args", return_value="invalid_key"):
+        with patch("httpx.AsyncClient") as mock_client:
+            mock_response_obj = Mock()
             mock_response_obj.json.return_value = mock_response
-            mock_response_obj.raise_for_status = AsyncMock()
-            mock_post.return_value = mock_response_obj
+            mock_response_obj.raise_for_status.return_value = None
+
+            mock_client_instance = AsyncMock()
+            mock_client_instance.post.return_value = mock_response_obj
+            mock_client.return_value.__aenter__.return_value = mock_client_instance
 
             result = await handle_pia_search({"query": "test"})
 
@@ -82,22 +78,112 @@ async def test_pia_search_api_error():
 @pytest.mark.asyncio
 async def test_pia_search_http_error():
     """Test PIA search with HTTP error."""
-    with patch.object(
-        type(settings),
-        "API_KEY",
-        new_callable=lambda: PropertyMock(return_value="test_key"),
-    ):
-        with patch("httpx.AsyncClient.post") as mock_post:
+    with patch.object(Settings, "_get_api_key_from_args", return_value="test_key"):
+        with patch("httpx.AsyncClient") as mock_client:
             mock_response_obj = AsyncMock()
             mock_response_obj.status_code = 500
             mock_response_obj.text = "Internal Server Error"
 
             http_error = httpx.HTTPStatusError(
                 "500 Server Error", request=AsyncMock(), response=mock_response_obj
             )
-            mock_post.side_effect = http_error
+
+            mock_client_instance = AsyncMock()
+            mock_client_instance.post.side_effect = http_error
+            mock_client.return_value.__aenter__.return_value = mock_client_instance
 
             result = await handle_pia_search({"query": "test"})
 
             assert len(result) == 1
             assert "HTTP Error 500" in result[0].text
+
+
+@pytest.mark.asyncio
+async def test_pia_search_facets_no_api_key():
+    """Test PIA search facets without API key."""
+    with patch.object(Settings, "_get_api_key_from_args", return_value=None):
+        result = await handle_pia_search_facets({"query": "test"})
+
+        assert len(result) == 1
+        assert "PIA API key is required" in result[0].text
+
+
+@pytest.mark.asyncio
+async def test_pia_search_facets_success():
+    """Test successful PIA search facets."""
+    mock_response = {
+        "jsonrpc": "2.0",
+        "id": 1,
+        "result": {
+            "facets": {
+                "data_source": ["OIG", "GAO", "CMS"],
+                "document_type": ["audit_report", "investigation", "guidance"],
+                "agency": ["HHS", "DOD", "VA"],
+            }
+        },
+    }
+
+    with patch.object(Settings, "_get_api_key_from_args", return_value="test_key"):
+        with patch("httpx.AsyncClient") as mock_client:
+            mock_response_obj = Mock()
+            mock_response_obj.json.return_value = mock_response
+            mock_response_obj.raise_for_status.return_value = None
+
+            mock_client_instance = AsyncMock()
+            mock_client_instance.post.return_value = mock_response_obj
+            mock_client.return_value.__aenter__.return_value = mock_client_instance
+
+            result = await handle_pia_search_facets({"query": "healthcare"})
+
+            assert len(result) == 1
+            assert "data_source" in result[0].text
+            assert "OIG" in result[0].text
+            assert "document_type" in result[0].text
+
+
+@pytest.mark.asyncio
+async def test_pia_search_facets_api_error():
+    """Test PIA search facets with API error."""
+    mock_response = {
+        "jsonrpc": "2.0",
+        "id": 1,
+        "error": {"code": -32000, "message": "Invalid query format"},
+    }
+
+    with patch.object(Settings, "_get_api_key_from_args", return_value="invalid_key"):
+        with patch("httpx.AsyncClient") as mock_client:
+            mock_response_obj = Mock()
+            mock_response_obj.json.return_value = mock_response
+            mock_response_obj.raise_for_status.return_value = None
+
+            mock_client_instance = AsyncMock()
+            mock_client_instance.post.return_value = mock_response_obj
+            mock_client.return_value.__aenter__.return_value = mock_client_instance
+
+            result = await handle_pia_search_facets({"query": "test"})
+
+            assert len(result) == 1
+            assert "API Error: Invalid query format" in result[0].text
+
+
+@pytest.mark.asyncio
+async def test_pia_search_facets_http_error():
+    """Test PIA search facets with HTTP error."""
+    with patch.object(Settings, "_get_api_key_from_args", return_value="test_key"):
+        with patch("httpx.AsyncClient") as mock_client:
+            mock_response_obj = AsyncMock()
+            mock_response_obj.status_code = 403
+            mock_response_obj.text = "Forbidden"
+
+            http_error = httpx.HTTPStatusError(
+                "403 Client Error", request=AsyncMock(), response=mock_response_obj
+            )
+
+            mock_client_instance = AsyncMock()
+            mock_client_instance.post.side_effect = http_error
+            mock_client.return_value.__aenter__.return_value = mock_client_instance
+
+            result = await handle_pia_search_facets({"query": "test"})
+
+            assert len(result) == 1
+            assert "HTTP Error 403" in result[0].text