Skip to content

Commit c457bb7

Browse files
committed
made fixes - doctests, pytest
Signed-off-by: Satya <[email protected]>
1 parent 947d626 commit c457bb7

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()
@@ -9254,7 +9250,7 @@ async def admin_list_a2a_agents(
92549250
92559251
Examples:
92569252
>>> import asyncio
9257-
>>> from unittest.mock import AsyncMock, MagicMock
9253+
>>> from unittest.mock import AsyncMock, MagicMock, patch
92589254
>>> from mcpgateway.schemas import A2AAgentRead, A2AAgentMetrics
92599255
>>> from datetime import datetime, timezone
92609256
>>>
@@ -9290,28 +9286,28 @@ async def admin_list_a2a_agents(
92909286
... )
92919287
... )
92929288
>>>
9293-
>>> original_list_agents_for_user = a2a_service.list_agents_for_user
9294-
>>> a2a_service.list_agents_for_user = AsyncMock(return_value=[mock_agent])
9295-
>>>
92969289
>>> async def test_admin_list_a2a_agents_active():
9297-
... result = await admin_list_a2a_agents(include_inactive=False, db=mock_db, user=mock_user)
9298-
... return len(result) > 0 and isinstance(result[0], dict) and result[0]['name'] == "Agent1"
9290+
... fake_service = MagicMock()
9291+
... fake_service.list_agents_for_user = AsyncMock(return_value=[mock_agent])
9292+
... with patch("mcpgateway.admin.a2a_service", new=fake_service):
9293+
... result = await admin_list_a2a_agents(include_inactive=False, db=mock_db, user=mock_user)
9294+
... return len(result) > 0 and isinstance(result[0], dict) and result[0]['name'] == "Agent1"
92999295
>>>
93009296
>>> asyncio.run(test_admin_list_a2a_agents_active())
93019297
True
93029298
>>>
9303-
>>> a2a_service.list_agents_for_user = AsyncMock(side_effect=Exception("A2A error"))
93049299
>>> async def test_admin_list_a2a_agents_exception():
9305-
... try:
9306-
... await admin_list_a2a_agents(False, db=mock_db, user=mock_user)
9307-
... return False
9308-
... except Exception as e:
9309-
... return "A2A error" in str(e)
9300+
... fake_service = MagicMock()
9301+
... fake_service.list_agents_for_user = AsyncMock(side_effect=Exception("A2A error"))
9302+
... with patch("mcpgateway.admin.a2a_service", new=fake_service):
9303+
... try:
9304+
... await admin_list_a2a_agents(False, db=mock_db, user=mock_user)
9305+
... return False
9306+
... except Exception as e:
9307+
... return "A2A error" in str(e)
93109308
>>>
93119309
>>> asyncio.run(test_admin_list_a2a_agents_exception())
93129310
True
9313-
>>>
9314-
>>> a2a_service.list_agents_for_user = original_list_agents_for_user
93159311
"""
93169312
if a2a_service is None:
93179313
LOGGER.warning("A2A features are disabled, returning empty list")
@@ -9565,48 +9561,81 @@ async def admin_edit_a2a_agent(
95659561
Returns:
95669562
JSONResponse: A JSON response indicating success or failure.
95679563
9568-
Example:
9569-
>>> import asyncio, json
9570-
>>> from unittest.mock import AsyncMock, MagicMock
9571-
>>> from fastapi import Request
9572-
>>> from fastapi.responses import JSONResponse
9573-
>>> from starlette.datastructures import FormData
9574-
>>> import types, mcpgateway.admin as admin_mod
9575-
>>> mock_db = MagicMock()
9576-
>>> mock_user = {"email": "[email protected]"}
9577-
>>> agent_id = "agent-123"
9578-
>>> form_data_success = FormData([
9579-
... ("name", "Updated Agent"),
9580-
... ("endpoint_url", "http://example.com/agent"),
9581-
... ("agent_type", "generic"),
9582-
... ("auth_type", "basic"),
9583-
... ("auth_username", "user"),
9584-
... ("auth_password", "pass"),
9585-
... ])
9586-
>>> mock_request_success = MagicMock(spec=Request, scope={"root_path": ""})
9587-
>>> mock_request_success.form = AsyncMock(return_value=form_data_success)
9588-
>>> admin_mod.get_user_email = lambda u: u["email"]
9589-
>>> admin_mod.get_oauth_encryption = lambda secret: types.SimpleNamespace(encrypt_secret=lambda s: f"enc({s})")
9590-
>>> admin_mod.settings = types.SimpleNamespace(auth_encryption_secret="dummy-secret")
9591-
>>> admin_mod.LOGGER = MagicMock()
9592-
>>> admin_mod.TeamManagementService = lambda db: types.SimpleNamespace(
9593-
... verify_team_for_user=AsyncMock(return_value="11111111-1111-1111-1111-111111111111")
9594-
... )
9595-
>>> admin_mod.MetadataCapture = types.SimpleNamespace(extract_modification_metadata=lambda req, u, _: {
9596-
... "modified_by": "[email protected]",
9597-
... "modified_from_ip": "127.0.0.1",
9598-
... "modified_via": "UI",
9599-
... "modified_user_agent": "pytest"
9600-
... })
9601-
>>> admin_mod.a2a_service = types.SimpleNamespace(update_agent=AsyncMock(return_value=True))
9602-
>>>
9603-
>>> async def test_admin_edit_a2a_agent_success():
9604-
... response = await admin_mod.admin_edit_a2a_agent(agent_id, mock_request_success, mock_db, mock_user)
9605-
... body = json.loads(response.body)
9606-
... return isinstance(response, JSONResponse) and response.status_code == 200 and body["success"] is True
9607-
>>>
9608-
>>> asyncio.run(test_admin_edit_a2a_agent_success())
9609-
True
9564+
Examples:
9565+
>>> import asyncio, json
9566+
>>> from unittest.mock import AsyncMock, MagicMock, patch
9567+
>>> from fastapi import Request
9568+
>>> from fastapi.responses import JSONResponse
9569+
>>> from starlette.datastructures import FormData
9570+
>>>
9571+
>>> mock_db = MagicMock()
9572+
>>> mock_user = {"email": "test_admin_user", "db": mock_db}
9573+
>>> agent_id = "agent-123"
9574+
>>>
9575+
>>> # Happy path: edit A2A agent successfully
9576+
>>> form_data_success = FormData([
9577+
... ("name", "Updated Agent"),
9578+
... ("endpoint_url", "http://updated-agent.com"),
9579+
... ("agent_type", "generic"),
9580+
... ("auth_type", "basic"),
9581+
... ("auth_username", "user"),
9582+
... ("auth_password", "pass"),
9583+
... ])
9584+
>>> mock_request_success = MagicMock(spec=Request, scope={"root_path": ""})
9585+
>>> mock_request_success.form = AsyncMock(return_value=form_data_success)
9586+
>>> original_update_agent = a2a_service.update_agent
9587+
>>> a2a_service.update_agent = AsyncMock()
9588+
>>>
9589+
>>> async def test_admin_edit_a2a_agent_success():
9590+
... response = await admin_edit_a2a_agent(agent_id, mock_request_success, mock_db, mock_user)
9591+
... body = json.loads(response.body)
9592+
... return isinstance(response, JSONResponse) and response.status_code == 200 and body["success"] is True
9593+
>>>
9594+
>>> asyncio.run(test_admin_edit_a2a_agent_success())
9595+
True
9596+
>>>
9597+
>>> # Error path: simulate exception during update
9598+
>>> form_data_error = FormData([
9599+
... ("name", "Error Agent"),
9600+
... ("endpoint_url", "http://error-agent.com"),
9601+
... ("auth_type", "basic"),
9602+
... ("auth_username", "user"),
9603+
... ("auth_password", "pass"),
9604+
... ])
9605+
>>> mock_request_error = MagicMock(spec=Request, scope={"root_path": ""})
9606+
>>> mock_request_error.form = AsyncMock(return_value=form_data_error)
9607+
>>> a2a_service.update_agent = AsyncMock(side_effect=Exception("Update failed"))
9608+
>>>
9609+
>>> async def test_admin_edit_a2a_agent_exception():
9610+
... response = await admin_edit_a2a_agent(agent_id, mock_request_error, mock_db, mock_user)
9611+
... body = json.loads(response.body)
9612+
... return isinstance(response, JSONResponse) and response.status_code == 500 and body["success"] is False and "Update failed" in body["message"]
9613+
>>>
9614+
>>> asyncio.run(test_admin_edit_a2a_agent_exception())
9615+
True
9616+
>>>
9617+
>>> # Validation error path: e.g., invalid URL
9618+
>>> form_data_validation = FormData([
9619+
... ("name", "Bad URL Agent"),
9620+
... ("endpoint_url", "invalid-url"),
9621+
... ("auth_type", "basic"),
9622+
... ("auth_username", "user"),
9623+
... ("auth_password", "pass"),
9624+
... ])
9625+
>>> mock_request_validation = MagicMock(spec=Request, scope={"root_path": ""})
9626+
>>> mock_request_validation.form = AsyncMock(return_value=form_data_validation)
9627+
>>>
9628+
>>> async def test_admin_edit_a2a_agent_validation():
9629+
... response = await admin_edit_a2a_agent(agent_id, mock_request_validation, mock_db, mock_user)
9630+
... body = json.loads(response.body)
9631+
... return isinstance(response, JSONResponse) and response.status_code in (422, 400) and body["success"] is False
9632+
>>>
9633+
>>> asyncio.run(test_admin_edit_a2a_agent_validation())
9634+
True
9635+
>>>
9636+
>>> # Restore original method
9637+
>>> a2a_service.update_agent = original_update_agent
9638+
96109639
"""
96119640

96129641
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)