diff --git a/server/secops/secops_mcp/tools/investigation_management.py b/server/secops/secops_mcp/tools/investigation_management.py index f419eee0..6737826d 100644 --- a/server/secops/secops_mcp/tools/investigation_management.py +++ b/server/secops/secops_mcp/tools/investigation_management.py @@ -76,19 +76,21 @@ async def list_investigations( """ try: chronicle = get_chronicle_client(project_id, customer_id, region) - print(f"Listing investigations (page_size={page_size})...") + logger.info(f"Listing investigations (page_size={page_size})...") result = chronicle.list_investigations( page_size=page_size, page_token=page_token ) investigations = result.get("investigations", []) - print(f"Successfully retrieved {len(investigations)} investigation(s)") + logger.info( + f"Successfully retrieved {len(investigations)} investigation(s)" + ) return result except Exception as e: error_msg = f"Error listing investigations: {str(e)}" - print(error_msg) + logger.error(error_msg, exc_info=True) return {"error": error_msg} @@ -149,7 +151,7 @@ async def get_investigation( } chronicle = get_chronicle_client(project_id, customer_id, region) - print(f"Retrieving investigation: {investigation_id}...") + logger.info(f"Retrieving investigation: {investigation_id}...") investigation = chronicle.get_investigation( investigation_id=investigation_id @@ -161,14 +163,16 @@ async def get_investigation( "investigation_id": investigation_id, } - print(f"Successfully retrieved investigation: {investigation_id}") + logger.info( + f"Successfully retrieved investigation: {investigation_id}" + ) return investigation except Exception as e: error_msg = ( f"Error retrieving investigation {investigation_id}: {str(e)}" ) - print(error_msg) + logger.error(error_msg, exc_info=True) return {"error": error_msg} @@ -226,7 +230,7 @@ async def trigger_investigation( } chronicle = get_chronicle_client(project_id, customer_id, region) - print(f"Triggering investigation for alert: {alert_id}...") + logger.info(f"Triggering investigation for alert: {alert_id}...") investigation = chronicle.trigger_investigation(alert_id=alert_id) @@ -250,14 +254,16 @@ async def trigger_investigation( }, } - print(f"Successfully triggered investigation for alert: {alert_id}") + logger.info( + f"Successfully triggered investigation for alert: {alert_id}" + ) return result except Exception as e: error_msg = ( f"Error triggering investigation for alert {alert_id}: {str(e)}" ) - print(error_msg) + logger.error(error_msg, exc_info=True) return {"error": error_msg} @@ -361,7 +367,7 @@ async def fetch_associated_investigations( detection_label = "alert" if is_alert_type else "case" ids = alert_ids if is_alert_type else case_ids - print( + logger.info( f"Fetching investigations for {len(ids)} " f"{detection_label}(s)..." ) @@ -409,7 +415,7 @@ async def fetch_associated_investigations( "associations": associations_dict, } - print( + logger.info( f"Successfully retrieved {total_investigations} " f"investigation(s) for {len(associations_dict)} " f"{detection_label}(s)" @@ -418,5 +424,5 @@ async def fetch_associated_investigations( except Exception as e: error_msg = f"Error fetching associated investigations: {str(e)}" - print(error_msg) + logger.error(error_msg, exc_info=True) return {"error": error_msg} diff --git a/server/secops/tests/test_investigation_management_unit.py b/server/secops/tests/test_investigation_management_unit.py new file mode 100644 index 00000000..ec2a3002 --- /dev/null +++ b/server/secops/tests/test_investigation_management_unit.py @@ -0,0 +1,225 @@ +# 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 investigation management diagnostics. + +The MCP stdio transport uses stdout for JSON-RPC, so these tools must report +progress and errors through the secops-mcp logger and leave stdout alone. +""" + +import logging +from unittest.mock import MagicMock, patch + +import pytest + +from secops_mcp.tools.investigation_management import ( + fetch_associated_investigations, + get_investigation, + list_investigations, + trigger_investigation, +) + +LOGGER_NAME = "secops-mcp" + + +@pytest.fixture +def chronicle_client(): + with patch( + "secops_mcp.tools.investigation_management.get_chronicle_client" + ) as factory: + client = MagicMock() + factory.return_value = client + yield client + + +def logs(caplog): + """Return (level, message) pairs emitted by the secops-mcp logger.""" + return [ + (record.levelname, record.getMessage()) + for record in caplog.records + if record.name == LOGGER_NAME + ] + + +@pytest.mark.asyncio +async def test_list_investigations_logs_progress( + chronicle_client, caplog, capsys +): + chronicle_client.list_investigations.return_value = { + "investigations": [{"name": "investigations/inv-1"}] + } + + with caplog.at_level(logging.INFO, logger=LOGGER_NAME): + result = await list_investigations(page_size=10) + + assert result == {"investigations": [{"name": "investigations/inv-1"}]} + assert capsys.readouterr().out == "" + assert logs(caplog) == [ + ("INFO", "Listing investigations (page_size=10)..."), + ("INFO", "Successfully retrieved 1 investigation(s)"), + ] + + +@pytest.mark.asyncio +async def test_list_investigations_logs_failure( + chronicle_client, caplog, capsys +): + chronicle_client.list_investigations.side_effect = RuntimeError("boom") + + with caplog.at_level(logging.INFO, logger=LOGGER_NAME): + result = await list_investigations() + + assert result == {"error": "Error listing investigations: boom"} + assert capsys.readouterr().out == "" + assert logs(caplog) == [ + ("INFO", "Listing investigations (page_size=50)..."), + ("ERROR", "Error listing investigations: boom"), + ] + assert caplog.records[-1].exc_info[0] is RuntimeError + + +@pytest.mark.asyncio +async def test_get_investigation_logs_progress( + chronicle_client, caplog, capsys +): + chronicle_client.get_investigation.return_value = { + "name": "investigations/inv-1" + } + + with caplog.at_level(logging.INFO, logger=LOGGER_NAME): + result = await get_investigation(investigation_id="inv-1") + + assert result == {"name": "investigations/inv-1"} + assert capsys.readouterr().out == "" + assert logs(caplog) == [ + ("INFO", "Retrieving investigation: inv-1..."), + ("INFO", "Successfully retrieved investigation: inv-1"), + ] + + +@pytest.mark.asyncio +async def test_get_investigation_logs_failure( + chronicle_client, caplog, capsys +): + chronicle_client.get_investigation.side_effect = ValueError("nope") + + with caplog.at_level(logging.INFO, logger=LOGGER_NAME): + result = await get_investigation(investigation_id="inv-1") + + assert result == {"error": "Error retrieving investigation inv-1: nope"} + assert capsys.readouterr().out == "" + assert logs(caplog) == [ + ("INFO", "Retrieving investigation: inv-1..."), + ("ERROR", "Error retrieving investigation inv-1: nope"), + ] + assert caplog.records[-1].exc_info[0] is ValueError + + +@pytest.mark.asyncio +async def test_trigger_investigation_logs_progress( + chronicle_client, caplog, capsys +): + chronicle_client.trigger_investigation.return_value = { + "name": "investigations/inv-1", + "displayName": "Triggered investigation", + "status": "RUNNING", + "triggerType": "MANUAL", + "createTime": "2026-05-28T18:58:18Z", + } + + with caplog.at_level(logging.INFO, logger=LOGGER_NAME): + result = await trigger_investigation(alert_id="alert-1") + + assert result["message"] == "Successfully triggered investigation" + assert capsys.readouterr().out == "" + assert logs(caplog) == [ + ("INFO", "Triggering investigation for alert: alert-1..."), + ("INFO", "Successfully triggered investigation for alert: alert-1"), + ] + + +@pytest.mark.asyncio +async def test_trigger_investigation_logs_failure( + chronicle_client, caplog, capsys +): + chronicle_client.trigger_investigation.side_effect = RuntimeError("down") + + with caplog.at_level(logging.INFO, logger=LOGGER_NAME): + result = await trigger_investigation(alert_id="alert-1") + + assert result == { + "error": "Error triggering investigation for alert alert-1: down" + } + assert capsys.readouterr().out == "" + assert logs(caplog) == [ + ("INFO", "Triggering investigation for alert: alert-1..."), + ("ERROR", "Error triggering investigation for alert alert-1: down"), + ] + assert caplog.records[-1].exc_info[0] is RuntimeError + + +@pytest.mark.asyncio +async def test_fetch_associated_investigations_logs_progress( + chronicle_client, caplog, capsys +): + chronicle_client.fetch_associated_investigations.return_value = { + "associationsList": { + "alert-1": { + "investigations": [ + { + "name": "investigations/inv-1", + "displayName": "Investigation 1", + "verdict": "MALICIOUS", + "confidence": "HIGH", + "status": "COMPLETE", + } + ] + } + } + } + + with caplog.at_level(logging.INFO, logger=LOGGER_NAME): + result = await fetch_associated_investigations( + detection_type="ALERT", alert_ids=["alert-1"] + ) + + assert result["total_investigations"] == 1 + assert capsys.readouterr().out == "" + assert logs(caplog) == [ + ("INFO", "Fetching investigations for 1 alert(s)..."), + ("INFO", "Successfully retrieved 1 investigation(s) for 1 alert(s)"), + ] + + +@pytest.mark.asyncio +async def test_fetch_associated_investigations_logs_failure( + chronicle_client, caplog, capsys +): + chronicle_client.fetch_associated_investigations.side_effect = ( + RuntimeError("denied") + ) + + with caplog.at_level(logging.INFO, logger=LOGGER_NAME): + result = await fetch_associated_investigations( + detection_type="CASE", case_ids=["case-1"] + ) + + assert result == { + "error": "Error fetching associated investigations: denied" + } + assert capsys.readouterr().out == "" + assert logs(caplog) == [ + ("INFO", "Fetching investigations for 1 case(s)..."), + ("ERROR", "Error fetching associated investigations: denied"), + ] + assert caplog.records[-1].exc_info[0] is RuntimeError