diff --git a/server/secops/secops_mcp/tools/security_alerts.py b/server/secops/secops_mcp/tools/security_alerts.py index b21afa08..912009a8 100644 --- a/server/secops/secops_mcp/tools/security_alerts.py +++ b/server/secops/secops_mcp/tools/security_alerts.py @@ -288,6 +288,6 @@ async def do_update_security_alert( chronicle = get_chronicle_client(project_id, customer_id, region) response = chronicle.update_alert(alert_id, reason=reason, status=status, verdict=verdict, comment=comment, root_cause=root_cause, priority=priority, severity=severity) except Exception as e: - return f'Error retrieving security alert for {alert_id}: {str(e)}' + return f'Error updating security alert for {alert_id}: {str(e)}' return json.dumps(response) diff --git a/server/secops/secops_mcp/tools/threat_intel.py b/server/secops/secops_mcp/tools/threat_intel.py index 8a47fe4a..e607701e 100644 --- a/server/secops/secops_mcp/tools/threat_intel.py +++ b/server/secops/secops_mcp/tools/threat_intel.py @@ -80,18 +80,20 @@ async def get_threat_intel( # Handle GeminiResponse object if hasattr(response, 'get_text_content'): # This is a GeminiResponse object, extract text content - return response.get_text_content() + text = response.get_text_content() + return text if isinstance(text, str) else json.dumps(text) elif hasattr(response, 'blocks') and isinstance(response.blocks, list): # Handle direct access to blocks if get_text_content isn't available text_content = [] for block in response.blocks: if hasattr(block, 'block_type') and hasattr(block, 'content'): if block.block_type == "TEXT": - text_content.append(block.content) + text_content.append(str(block.content)) return "\n\n".join(text_content) if text_content else "No text content found in response." elif isinstance(response, dict) and 'answer' in response: # Legacy format or different API response - return response.get('answer', 'No answer was provided by the model.') + answer = response.get('answer', 'No answer was provided by the model.') + return answer if isinstance(answer, str) else json.dumps(answer) elif isinstance(response, str): # Direct string response return response diff --git a/server/secops/tests/test_secops_mcp.py b/server/secops/tests/test_secops_mcp.py index 05ef0fbf..a5e5d636 100644 --- a/server/secops/tests/test_secops_mcp.py +++ b/server/secops/tests/test_secops_mcp.py @@ -14,6 +14,7 @@ 3. Run: pytest -xvs server/secops/tests/test_secops_mcp.py """ +import json import os import uuid from typing import Dict @@ -476,8 +477,9 @@ async def test_do_update_security_alert(self, chronicle_config: Dict[str, str]) alert_id=alert_id ) - # This should return either a valid response or an error dict - assert isinstance(result, dict) + # This should return a JSON-encoded string representing the response + assert isinstance(result, str) + assert isinstance(json.loads(result), dict) @pytest.mark.asyncio async def test_get_security_alert_by_id(self, chronicle_config: Dict[str, str]) -> None: @@ -498,8 +500,9 @@ async def test_get_security_alert_by_id(self, chronicle_config: Dict[str, str]) alert_id=alert_id ) - # This should return either a valid response or an error dict - assert isinstance(result, dict) + # This should return a JSON-encoded string representing the response + assert isinstance(result, str) + assert isinstance(json.loads(result), dict) @pytest.mark.asyncio async def test_list_feeds(self, chronicle_config: Dict[str, str]) -> None: diff --git a/server/secops/tests/test_security_alerts_unit.py b/server/secops/tests/test_security_alerts_unit.py index e2cc0140..6b67c2fd 100644 --- a/server/secops/tests/test_security_alerts_unit.py +++ b/server/secops/tests/test_security_alerts_unit.py @@ -17,8 +17,11 @@ from unittest.mock import MagicMock, patch import pytest - -from secops_mcp.tools.security_alerts import get_security_alerts +from secops_mcp.tools.security_alerts import ( + do_update_security_alert, + get_security_alert_by_id, + get_security_alerts, +) @pytest.fixture @@ -131,3 +134,78 @@ async def test_get_security_alerts_preserves_unspecified_verdict(chronicle_clien ) assert "Verdict: VERDICT_UNSPECIFIED" in output + + +@pytest.mark.asyncio +async def test_get_security_alerts_empty_returns_str(chronicle_client): + chronicle_client.get_alerts.return_value = [] + + result = await get_security_alerts(project_id="test", customer_id="test") + + assert isinstance(result, str) + assert result == "No security alerts found for the specified time range." + + +@pytest.mark.asyncio +async def test_get_security_alerts_error_returns_str(chronicle_client): + chronicle_client.get_alerts.side_effect = RuntimeError("Chronicle error") + + result = await get_security_alerts(project_id="test", customer_id="test") + + assert isinstance(result, str) + assert "Error retrieving security alerts: Chronicle error" in result + + +@pytest.mark.asyncio +async def test_get_security_alert_by_id_returns_json_string(chronicle_client): + mock_alert = { + "id": "de_f47e71ca", + "detection": [{"ruleName": "Phishing"}], + "createdTime": "2026-05-28T18:58:18Z", + "status": "OPEN", + } + chronicle_client.get_alert.return_value = mock_alert + + result = await get_security_alert_by_id(alert_id="de_f47e71ca") + + assert isinstance(result, str) + parsed = json.loads(result) + assert parsed == mock_alert + + +@pytest.mark.asyncio +async def test_get_security_alert_by_id_handles_error(chronicle_client): + chronicle_client.get_alert.side_effect = Exception("Alert not found") + + result = await get_security_alert_by_id(alert_id="de_invalid") + + assert isinstance(result, str) + assert "Error retrieving security alert for de_invalid: Alert not found" in result + + +@pytest.mark.asyncio +async def test_do_update_security_alert_returns_json_string(chronicle_client): + mock_update = { + "id": "de_f47e71ca", + "status": "CLOSED", + "verdict": "FALSE_POSITIVE", + } + chronicle_client.update_alert.return_value = mock_update + + result = await do_update_security_alert( + alert_id="de_f47e71ca", status="CLOSED", verdict="FALSE_POSITIVE" + ) + + assert isinstance(result, str) + parsed = json.loads(result) + assert parsed == mock_update + + +@pytest.mark.asyncio +async def test_do_update_security_alert_handles_error(chronicle_client): + chronicle_client.update_alert.side_effect = Exception("Update failed") + + result = await do_update_security_alert(alert_id="de_f47e71ca", status="CLOSED") + + assert isinstance(result, str) + assert result == "Error updating security alert for de_f47e71ca: Update failed" diff --git a/server/secops/tests/test_threat_intel_unit.py b/server/secops/tests/test_threat_intel_unit.py new file mode 100644 index 00000000..b6a3e131 --- /dev/null +++ b/server/secops/tests/test_threat_intel_unit.py @@ -0,0 +1,118 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for threat intelligence tool return types and serialization.""" + +import json +from unittest.mock import MagicMock, patch + +import pytest +from secops_mcp.tools.threat_intel import get_threat_intel + + +@pytest.fixture +def chronicle_client(): + with patch("secops_mcp.tools.threat_intel.get_chronicle_client") as mock_get_client: + client = MagicMock() + mock_get_client.return_value = client + yield client + + +@pytest.mark.asyncio +async def test_get_threat_intel_gemini_response_object(chronicle_client): + mock_resp = MagicMock() + mock_resp.get_text_content.return_value = "APT41 is a prolific cyber threat group." + chronicle_client.gemini.return_value = mock_resp + + result = await get_threat_intel(query="Summarize APT41") + + assert isinstance(result, str) + assert result == "APT41 is a prolific cyber threat group." + + +@pytest.mark.asyncio +async def test_get_threat_intel_blocks_format(chronicle_client): + block1 = MagicMock() + block1.block_type = "TEXT" + block1.content = "Paragraph 1" + + block2 = MagicMock() + block2.block_type = "TEXT" + block2.content = "Paragraph 2" + + mock_resp = MagicMock(spec=["blocks"]) + mock_resp.blocks = [block1, block2] + chronicle_client.gemini.return_value = mock_resp + + result = await get_threat_intel(query="Test query") + + assert isinstance(result, str) + assert result == "Paragraph 1\n\nParagraph 2" + + +@pytest.mark.asyncio +async def test_get_threat_intel_dict_with_str_answer(chronicle_client): + chronicle_client.gemini.return_value = {"answer": "Threat intel summary."} + + result = await get_threat_intel(query="Test query") + + assert isinstance(result, str) + assert result == "Threat intel summary." + + +@pytest.mark.asyncio +async def test_get_threat_intel_dict_with_nested_dict_answer(chronicle_client): + chronicle_client.gemini.return_value = { + "answer": {"summary": "Nested dict content", "risk_level": "High"} + } + + result = await get_threat_intel(query="Test query") + + assert isinstance(result, str) + parsed = json.loads(result) + assert parsed == {"summary": "Nested dict content", "risk_level": "High"} + + +@pytest.mark.asyncio +async def test_get_threat_intel_direct_str(chronicle_client): + chronicle_client.gemini.return_value = "Direct string response" + + result = await get_threat_intel(query="Test query") + + assert isinstance(result, str) + assert result == "Direct string response" + + +@pytest.mark.asyncio +async def test_get_threat_intel_unexpected_format_serializes_to_json(chronicle_client): + unexpected_response = { + "results": [{"actor": "APT29", "confidence": 95}], + "metadata": {"total": 1}, + } + chronicle_client.gemini.return_value = unexpected_response + + result = await get_threat_intel(query="Test query") + + assert isinstance(result, str) + parsed = json.loads(result) + assert parsed == unexpected_response + + +@pytest.mark.asyncio +async def test_get_threat_intel_handles_exception(chronicle_client): + chronicle_client.gemini.side_effect = RuntimeError("Chronicle API unavailable") + + result = await get_threat_intel(query="Test query") + + assert isinstance(result, str) + assert "Error retrieving threat intelligence: Chronicle API unavailable" in result