44import time
55from io import StringIO
66from pathlib import Path
7+ from typing import Optional
78from unittest .mock import ANY , MagicMock , patch
89
910import pytest
@@ -743,26 +744,22 @@ def _guardrail_config_with_mcp_server(agents: dict) -> GuardrailConfig:
743744 ],
744745)
745746def test_mcp_server_statuses_not_fetched_while_the_guardrail_is_off (
746- config : object , monkeypatch : pytest . MonkeyPatch , mock_save_mcp_server_statuses : MagicMock
747+ config : Optional [ GuardrailConfig ] , mock_save_mcp_server_statuses : MagicMock
747748) -> None :
748- monkeypatch .setattr (_session_start_mod , 'load_guardrail_config' , MagicMock (return_value = config ))
749749 ai_client = MagicMock ()
750750
751- _session_start_mod ._sync_mcp_server_statuses (ai_client , 'tenant-a' , force = True )
751+ _session_start_mod ._sync_mcp_server_statuses (ai_client , config , 'tenant-a' , force = True )
752752
753753 ai_client .get_mcp_server_statuses .assert_not_called ()
754754 mock_save_mcp_server_statuses .assert_not_called ()
755755
756756
757- def test_mcp_server_statuses_fetched_when_on_for_some_agent (
758- monkeypatch : pytest .MonkeyPatch , mock_save_mcp_server_statuses : MagicMock
759- ) -> None :
757+ def test_mcp_server_statuses_fetched_when_on_for_some_agent (mock_save_mcp_server_statuses : MagicMock ) -> None :
760758 config = _guardrail_config_with_mcp_server ({'cursor' : 'Off' , 'claude-code' : 'Report' })
761- monkeypatch .setattr (_session_start_mod , 'load_guardrail_config' , MagicMock (return_value = config ))
762759 ai_client = MagicMock ()
763760 ai_client .get_mcp_server_statuses .return_value = _STATUSES_RESPONSE
764761
765- _session_start_mod ._sync_mcp_server_statuses (ai_client , 'tenant-a' , force = False )
762+ _session_start_mod ._sync_mcp_server_statuses (ai_client , config , 'tenant-a' , force = False )
766763
767764 # Cached with the guardrail config's TTL and the policy's strict-mode flag.
768765 mock_save_mcp_server_statuses .assert_called_once_with (_SERVERS , 'tenant-a' , 600 , True )
@@ -773,32 +770,28 @@ def test_fresh_mcp_server_statuses_skip_the_fetch_unless_forced(
773770) -> None :
774771 config = _guardrail_config_with_mcp_server ({'cursor' : 'Block' })
775772 cached = McpServerStatuses (servers = [], fetched_at = time .time (), tenant_id = 'tenant-a' )
776- monkeypatch .setattr (_session_start_mod , 'load_guardrail_config' , MagicMock (return_value = config ))
777773 monkeypatch .setattr (_session_start_mod , 'load_mcp_server_statuses' , MagicMock (return_value = cached ))
778774 ai_client = MagicMock ()
779775 ai_client .get_mcp_server_statuses .return_value = _STATUSES_RESPONSE
780776
781- _session_start_mod ._sync_mcp_server_statuses (ai_client , 'tenant-a' , force = False )
777+ _session_start_mod ._sync_mcp_server_statuses (ai_client , config , 'tenant-a' , force = False )
782778 ai_client .get_mcp_server_statuses .assert_not_called ()
783779
784780 # Another tenant's cache is refetched.
785- _session_start_mod ._sync_mcp_server_statuses (ai_client , 'tenant-b' , force = False )
781+ _session_start_mod ._sync_mcp_server_statuses (ai_client , config , 'tenant-b' , force = False )
786782 assert ai_client .get_mcp_server_statuses .call_count == 1
787783
788784 # A just-reported session context forces a refetch.
789- _session_start_mod ._sync_mcp_server_statuses (ai_client , 'tenant-a' , force = True )
785+ _session_start_mod ._sync_mcp_server_statuses (ai_client , config , 'tenant-a' , force = True )
790786 assert ai_client .get_mcp_server_statuses .call_count == 2
791787
792788
793- def test_failed_mcp_server_statuses_fetch_keeps_the_cache (
794- monkeypatch : pytest .MonkeyPatch , mock_save_mcp_server_statuses : MagicMock
795- ) -> None :
789+ def test_failed_mcp_server_statuses_fetch_keeps_the_cache (mock_save_mcp_server_statuses : MagicMock ) -> None :
796790 config = _guardrail_config_with_mcp_server ({'cursor' : 'Block' })
797- monkeypatch .setattr (_session_start_mod , 'load_guardrail_config' , MagicMock (return_value = config ))
798791 ai_client = MagicMock ()
799792 ai_client .get_mcp_server_statuses .return_value = None
800793
801- _session_start_mod ._sync_mcp_server_statuses (ai_client , 'tenant-a' , force = True )
794+ _session_start_mod ._sync_mcp_server_statuses (ai_client , config , 'tenant-a' , force = True )
802795
803796 mock_save_mcp_server_statuses .assert_not_called ()
804797
0 commit comments