Skip to content

Commit b13c988

Browse files
TS0713TS0713
committed
made fixes - doctests, pytest
Signed-off-by: Satya <[email protected]>
1 parent b14c246 commit b13c988

File tree

3 files changed

+124
-84
lines changed

3 files changed

+124
-84
lines changed

‎mcpgateway/admin.py‎

Lines changed: 91 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -5349,7 +5349,6 @@ async def admin_edit_tool(
53495349
an error message if the update fails.
53505350
53515351
Examples:
5352-
Examples:
53535352
>>> import asyncio
53545353
>>> from unittest.mock import AsyncMock, MagicMock
53555354
>>> from fastapi import Request
@@ -5484,7 +5483,6 @@ async def admin_edit_tool(
54845483
54855484
>>> # Restore original method
54865485
>>> tool_service.update_tool = original_update_tool
5487-
54885486
"""
54895487
LOGGER.debug(f"User {get_user_email(user)} is editing tool ID {tool_id}")
54905488
form = await request.form()
@@ -6249,10 +6247,7 @@ async def admin_edit_gateway(
62496247
>>>
62506248
>>> async def test_admin_edit_gateway_success():
62516249
... response = await admin_edit_gateway(gateway_id, mock_request_success, mock_db, mock_user)
6252-
... #print("Response:", response)
6253-
... #print("Status code:", response.status_code)
6254-
... #print("Body:", response.body)
6255-
... return isinstance(response, JSONResponse) and response.status_code == 200 and json.loads(response.body)["success"]
6250+
... return isinstance(response, JSONResponse) and response.status_code == 200 and json.loads(response.body)["success"] is True
62566251
>>>
62576252
>>> asyncio.run(test_admin_edit_gateway_success())
62586253
True
@@ -7414,6 +7409,7 @@ async def admin_edit_prompt(
74147409
>>> asyncio.run(test_admin_edit_prompt_inactive())
74157410
True
74167411
>>> prompt_service.update_prompt = original_update_prompt
7412+
74177413
"""
74187414
LOGGER.debug(f"User {get_user_email(user)} is editing prompt {prompt_id}")
74197415
form = await request.form()
@@ -9251,7 +9247,7 @@ async def admin_list_a2a_agents(
92519247
92529248
Examples:
92539249
>>> import asyncio
9254-
>>> from unittest.mock import AsyncMock, MagicMock
9250+
>>> from unittest.mock import AsyncMock, MagicMock, patch
92559251
>>> from mcpgateway.schemas import A2AAgentRead, A2AAgentMetrics
92569252
>>> from datetime import datetime, timezone
92579253
>>>
@@ -9287,28 +9283,28 @@ async def admin_list_a2a_agents(
92879283
... )
92889284
... )
92899285
>>>
9290-
>>> original_list_agents_for_user = a2a_service.list_agents_for_user
9291-
>>> a2a_service.list_agents_for_user = AsyncMock(return_value=[mock_agent])
9292-
>>>
92939286
>>> async def test_admin_list_a2a_agents_active():
9294-
... result = await admin_list_a2a_agents(include_inactive=False, db=mock_db, user=mock_user)
9295-
... return len(result) > 0 and isinstance(result[0], dict) and result[0]['name'] == "Agent1"
9287+
... fake_service = MagicMock()
9288+
... fake_service.list_agents_for_user = AsyncMock(return_value=[mock_agent])
9289+
... with patch("mcpgateway.admin.a2a_service", new=fake_service):
9290+
... result = await admin_list_a2a_agents(include_inactive=False, db=mock_db, user=mock_user)
9291+
... return len(result) > 0 and isinstance(result[0], dict) and result[0]['name'] == "Agent1"
92969292
>>>
92979293
>>> asyncio.run(test_admin_list_a2a_agents_active())
92989294
True
92999295
>>>
9300-
>>> a2a_service.list_agents_for_user = AsyncMock(side_effect=Exception("A2A error"))
93019296
>>> async def test_admin_list_a2a_agents_exception():
9302-
... try:
9303-
... await admin_list_a2a_agents(False, db=mock_db, user=mock_user)
9304-
... return False
9305-
... except Exception as e:
9306-
... return "A2A error" in str(e)
9297+
... fake_service = MagicMock()
9298+
... fake_service.list_agents_for_user = AsyncMock(side_effect=Exception("A2A error"))
9299+
... with patch("mcpgateway.admin.a2a_service", new=fake_service):
9300+
... try:
9301+
... await admin_list_a2a_agents(False, db=mock_db, user=mock_user)
9302+
... return False
9303+
... except Exception as e:
9304+
... return "A2A error" in str(e)
93079305
>>>
93089306
>>> asyncio.run(test_admin_list_a2a_agents_exception())
93099307
True
9310-
>>>
9311-
>>> a2a_service.list_agents_for_user = original_list_agents_for_user
93129308
"""
93139309
if a2a_service is None:
93149310
LOGGER.warning("A2A features are disabled, returning empty list")
@@ -9562,48 +9558,81 @@ async def admin_edit_a2a_agent(
95629558
Returns:
95639559
JSONResponse: A JSON response indicating success or failure.
95649560
9565-
Example:
9566-
>>> import asyncio, json
9567-
>>> from unittest.mock import AsyncMock, MagicMock
9568-
>>> from fastapi import Request
9569-
>>> from fastapi.responses import JSONResponse
9570-
>>> from starlette.datastructures import FormData
9571-
>>> import types, mcpgateway.admin as admin_mod
9572-
>>> mock_db = MagicMock()
9573-
>>> mock_user = {"email": "[email protected]"}
9574-
>>> agent_id = "agent-123"
9575-
>>> form_data_success = FormData([
9576-
... ("name", "Updated Agent"),
9577-
... ("endpoint_url", "http://example.com/agent"),
9578-
... ("agent_type", "generic"),
9579-
... ("auth_type", "basic"),
9580-
... ("auth_username", "user"),
9581-
... ("auth_password", "pass"),
9582-
... ])
9583-
>>> mock_request_success = MagicMock(spec=Request, scope={"root_path": ""})
9584-
>>> mock_request_success.form = AsyncMock(return_value=form_data_success)
9585-
>>> admin_mod.get_user_email = lambda u: u["email"]
9586-
>>> admin_mod.get_oauth_encryption = lambda secret: types.SimpleNamespace(encrypt_secret=lambda s: f"enc({s})")
9587-
>>> admin_mod.settings = types.SimpleNamespace(auth_encryption_secret="dummy-secret")
9588-
>>> admin_mod.LOGGER = MagicMock()
9589-
>>> admin_mod.TeamManagementService = lambda db: types.SimpleNamespace(
9590-
... verify_team_for_user=AsyncMock(return_value="11111111-1111-1111-1111-111111111111")
9591-
... )
9592-
>>> admin_mod.MetadataCapture = types.SimpleNamespace(extract_modification_metadata=lambda req, u, _: {
9593-
... "modified_by": "[email protected]",
9594-
... "modified_from_ip": "127.0.0.1",
9595-
... "modified_via": "UI",
9596-
... "modified_user_agent": "pytest"
9597-
... })
9598-
>>> admin_mod.a2a_service = types.SimpleNamespace(update_agent=AsyncMock(return_value=True))
9599-
>>>
9600-
>>> async def test_admin_edit_a2a_agent_success():
9601-
... response = await admin_mod.admin_edit_a2a_agent(agent_id, mock_request_success, mock_db, mock_user)
9602-
... body = json.loads(response.body)
9603-
... return isinstance(response, JSONResponse) and response.status_code == 200 and body["success"] is True
9604-
>>>
9605-
>>> asyncio.run(test_admin_edit_a2a_agent_success())
9606-
True
9561+
Examples:
9562+
>>> import asyncio, json
9563+
>>> from unittest.mock import AsyncMock, MagicMock, patch
9564+
>>> from fastapi import Request
9565+
>>> from fastapi.responses import JSONResponse
9566+
>>> from starlette.datastructures import FormData
9567+
>>>
9568+
>>> mock_db = MagicMock()
9569+
>>> mock_user = {"email": "test_admin_user", "db": mock_db}
9570+
>>> agent_id = "agent-123"
9571+
>>>
9572+
>>> # Happy path: edit A2A agent successfully
9573+
>>> form_data_success = FormData([
9574+
... ("name", "Updated Agent"),
9575+
... ("endpoint_url", "http://updated-agent.com"),
9576+
... ("agent_type", "generic"),
9577+
... ("auth_type", "basic"),
9578+
... ("auth_username", "user"),
9579+
... ("auth_password", "pass"),
9580+
... ])
9581+
>>> mock_request_success = MagicMock(spec=Request, scope={"root_path": ""})
9582+
>>> mock_request_success.form = AsyncMock(return_value=form_data_success)
9583+
>>> original_update_agent = a2a_service.update_agent
9584+
>>> a2a_service.update_agent = AsyncMock()
9585+
>>>
9586+
>>> async def test_admin_edit_a2a_agent_success():
9587+
... response = await admin_edit_a2a_agent(agent_id, mock_request_success, mock_db, mock_user)
9588+
... body = json.loads(response.body)
9589+
... return isinstance(response, JSONResponse) and response.status_code == 200 and body["success"] is True
9590+
>>>
9591+
>>> asyncio.run(test_admin_edit_a2a_agent_success())
9592+
True
9593+
>>>
9594+
>>> # Error path: simulate exception during update
9595+
>>> form_data_error = FormData([
9596+
... ("name", "Error Agent"),
9597+
... ("endpoint_url", "http://error-agent.com"),
9598+
... ("auth_type", "basic"),
9599+
... ("auth_username", "user"),
9600+
... ("auth_password", "pass"),
9601+
... ])
9602+
>>> mock_request_error = MagicMock(spec=Request, scope={"root_path": ""})
9603+
>>> mock_request_error.form = AsyncMock(return_value=form_data_error)
9604+
>>> a2a_service.update_agent = AsyncMock(side_effect=Exception("Update failed"))
9605+
>>>
9606+
>>> async def test_admin_edit_a2a_agent_exception():
9607+
... response = await admin_edit_a2a_agent(agent_id, mock_request_error, mock_db, mock_user)
9608+
... body = json.loads(response.body)
9609+
... return isinstance(response, JSONResponse) and response.status_code == 500 and body["success"] is False and "Update failed" in body["message"]
9610+
>>>
9611+
>>> asyncio.run(test_admin_edit_a2a_agent_exception())
9612+
True
9613+
>>>
9614+
>>> # Validation error path: e.g., invalid URL
9615+
>>> form_data_validation = FormData([
9616+
... ("name", "Bad URL Agent"),
9617+
... ("endpoint_url", "invalid-url"),
9618+
... ("auth_type", "basic"),
9619+
... ("auth_username", "user"),
9620+
... ("auth_password", "pass"),
9621+
... ])
9622+
>>> mock_request_validation = MagicMock(spec=Request, scope={"root_path": ""})
9623+
>>> mock_request_validation.form = AsyncMock(return_value=form_data_validation)
9624+
>>>
9625+
>>> async def test_admin_edit_a2a_agent_validation():
9626+
... response = await admin_edit_a2a_agent(agent_id, mock_request_validation, mock_db, mock_user)
9627+
... body = json.loads(response.body)
9628+
... return isinstance(response, JSONResponse) and response.status_code in (422, 400) and body["success"] is False
9629+
>>>
9630+
>>> asyncio.run(test_admin_edit_a2a_agent_validation())
9631+
True
9632+
>>>
9633+
>>> # Restore original method
9634+
>>> a2a_service.update_agent = original_update_agent
9635+
96079636
"""
96089637

96099638
try:

‎mcpgateway/cache/resource_cache.py‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -125,18 +125,24 @@ def get(self, key: str) -> Optional[Any]:
125125
126126
Examples:
127127
>>> from mcpgateway.cache.resource_cache import ResourceCache
128+
>>> from unittest.mock import patch
129+
130+
>>> # Normal get
128131
>>> cache = ResourceCache(max_size=2, ttl=1)
129132
>>> cache.set('a', 1)
130133
>>> cache.get('a')
131134
1
132-
>>> # Test expiration by using a very short TTL
133-
>>> short_cache = ResourceCache(max_size=2, ttl=0.1)
134-
>>> short_cache.set('b', 2)
135-
>>> short_cache.get('b')
135+
136+
>>> # Test expiration deterministically using mock time
137+
>>> with patch("time.time") as mock_time:
138+
... mock_time.return_value = 1000
139+
... short_cache = ResourceCache(max_size=2, ttl=0.1)
140+
... short_cache.set('b', 2)
141+
... short_cache.get('b')
142+
... # Advance time past TTL
143+
... mock_time.return_value = 1000.2
144+
... short_cache.get('b') is None
136145
2
137-
>>> import time
138-
>>> time.sleep(0.2) # Sleep longer than TTL (0.1s) to ensure expiration
139-
>>> short_cache.get('b') is None
140146
True
141147
"""
142148
if key not in self._cache:

‎tests/unit/mcpgateway/services/test_a2a_service.py‎

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
from mcpgateway.schemas import A2AAgentCreate, A2AAgentUpdate
2222
from mcpgateway.services.a2a_service import A2AAgentError, A2AAgentNameConflictError, A2AAgentNotFoundError, A2AAgentService
2323

24-
2524
class TestA2AAgentService:
2625
"""Test suite for A2A Agent Service."""
2726

@@ -43,11 +42,13 @@ def sample_agent_create(self):
4342
description="Test agent for unit tests",
4443
endpoint_url="https://api.example.com/agent",
4544
agent_type="custom",
45+
auth_username="user",
46+
auth_password="dummy_pass",
4647
protocol_version="1.0",
4748
capabilities={"chat": True, "tools": False},
4849
config={"max_tokens": 1000},
49-
auth_type="api_key",
50-
auth_value="test-api-key",
50+
auth_type="basic",
51+
auth_value="encode-auth-value",
5152
tags=["test", "ai"],
5253
)
5354

@@ -65,8 +66,8 @@ def sample_db_agent(self):
6566
protocol_version="1.0",
6667
capabilities={"chat": True, "tools": False},
6768
config={"max_tokens": 1000},
68-
auth_type="api_key",
69-
auth_value="test-api-key",
69+
auth_type="basic",
70+
auth_value="encoded-auth-value",
7071
enabled=True,
7172
reachable=True,
7273
tags=["test", "ai"],
@@ -405,14 +406,11 @@ async def test_reset_metrics_specific_agent(self, service, mock_db):
405406
mock_db.commit.assert_called_once()
406407

407408
def test_db_to_schema_conversion(self, service, sample_db_agent):
408-
"""Test database model to schema conversion with db parameter."""
409-
from unittest.mock import MagicMock
410-
from datetime import datetime, timezone
409+
"""
410+
Test database model to schema conversion with db parameter.
411+
"""
411412

412-
# Create a mock DB session
413413
mock_db = MagicMock()
414-
415-
# Mock _get_team_name to return a test team name
416414
service._get_team_name = MagicMock(return_value="Test Team")
417415

418416
# Add some mock metrics
@@ -428,7 +426,10 @@ def test_db_to_schema_conversion(self, service, sample_db_agent):
428426

429427
sample_db_agent.metrics = [metric1, metric2]
430428

431-
# Set all required attributes that might be missing
429+
# Add dummy auth_value (doesn't matter since we'll patch decode_auth)
430+
sample_db_agent.auth_value = "fake_encrypted_auth"
431+
432+
# Set all required attributes
432433
sample_db_agent.created_by = "test_user"
433434
sample_db_agent.created_from_ip = "127.0.0.1"
434435
sample_db_agent.created_via = "test"
@@ -441,9 +442,13 @@ def test_db_to_schema_conversion(self, service, sample_db_agent):
441442
sample_db_agent.federation_source = None
442443
sample_db_agent.version = 1
443444
sample_db_agent.visibility = "private"
445+
sample_db_agent.auth_type = "none"
446+
sample_db_agent.auth_header_key = "Authorization"
447+
sample_db_agent.auth_header_value = "Basic dGVzdDp2YWx1ZQ==" # base64 for "test:value"
444448

445-
# Execute with db parameter
446-
result = service._db_to_schema(mock_db, sample_db_agent)
449+
# Patch decode_auth to return a dummy decoded dict
450+
with patch("mcpgateway.schemas.decode_auth", return_value={"user": "decoded"}):
451+
result = service._db_to_schema(mock_db, sample_db_agent)
447452

448453
# Verify
449454
assert result.id == sample_db_agent.id
@@ -453,7 +458,7 @@ def test_db_to_schema_conversion(self, service, sample_db_agent):
453458
assert result.metrics.failed_executions == 1
454459
assert result.metrics.failure_rate == 50.0
455460
assert result.metrics.avg_response_time == 1.5
456-
assert result.team == "Test Team" # Check that the mocked team name is set
461+
assert result.team == "Test Team"
457462

458463

459464
class TestA2AAgentIntegration:

0 commit comments

Comments
 (0)