PyTorch HUD MCP Server
by izaitsevfb
- test
#!/usr/bin/env python3
"""
Comprehensive test for async MCP resource endpoints
This test script ensures that all async functions in the PyTorch HUD API are properly
exposed via resource endpoints in the MCP server and that these endpoints correctly
await the underlying async functions.
"""
import json
import unittest
from unittest.mock import patch
# Import the resource endpoints from MCP server
from pytorch_hud.server.mcp_server import (
get_job_details_resource, get_recent_commits_with_jobs_resource
)
class TestAsyncMCPEndpoints(unittest.IsolatedAsyncioTestCase):
"""Tests for the async MCP resource endpoints."""
@patch('pytorch_hud.server.mcp_server.get_recent_commits_with_jobs')
async def test_recent_commits_with_jobs_resource(self, mock_get_recent_commits):
"""Test that the universal resource endpoint properly awaits the async function."""
# Create mock return value that contains filtered jobs
mock_result = {
"repo": "pytorch/pytorch",
"branch_or_commit_sha": "main",
"commits": [
{
"sha": "abcdef",
"title": "Test Commit",
"author": "test-user",
"status": "red",
"job_counts": {
"total": 5,
"success": 3,
"failure": 2
},
"jobs": [
{"id": "job1", "name": "test_job1", "conclusion": "failure"},
{"id": "job2", "name": "test_job2", "conclusion": "failure"}
]
}
],
"pagination": {
"page": 1,
"per_page": 10,
"total_commits": 1,
"returned_commits": 1
},
"filters": {
"include_success": False,
"include_pending": False,
"include_failures": True,
"job_name_filter_regex": None
},
"_metadata": {
"timestamp": "2025-03-07T12:00:00Z"
}
}
# Set up mock to return our sample data
mock_get_recent_commits.return_value = mock_result
# Call the universal resource endpoint with failure filter
result = await get_recent_commits_with_jobs_resource(
"pytorch", "pytorch", "main",
include_success=False,
include_pending=False,
include_failures=True,
page=1,
per_page=10
)
# Verify the resource endpoint was called with the correct arguments
mock_get_recent_commits.assert_called_once_with(
repo_owner="pytorch",
repo_name="pytorch",
branch_or_commit_sha="main",
include_success=False,
include_pending=False,
include_failures=True,
include_commit_details=True,
job_name_filter_regex=None,
failure_line_filter_regex=None,
page=1,
per_page=10,
ctx=None
)
# Result should be a JSON string - parse it back to verify contents
result_data = json.loads(result)
# Check that it contains the expected fields for filtered jobs
self.assertIn("commits", result_data)
self.assertEqual(len(result_data["commits"]), 1)
self.assertEqual(len(result_data["commits"][0]["jobs"]), 2)
# Reset mock for next test
mock_get_recent_commits.reset_mock()
# Create mock return value that contains only commit data (no jobs)
mock_result_commits = {
"repo": "pytorch/pytorch",
"branch_or_commit_sha": "main",
"commits": [
{
"sha": "commit1",
"title": "Test Commit 1",
"author": "test-user",
"status": "green",
"job_counts": {"total": 10, "success": 10}
},
{
"sha": "commit2",
"title": "Test Commit 2",
"author": "test-user",
"status": "red",
"job_counts": {"total": 8, "success": 5, "failure": 3}
}
],
"pagination": {
"page": 1,
"per_page": 10,
"total_commits": 2,
"returned_commits": 2
},
"_metadata": {
"timestamp": "2025-03-07T12:00:00Z"
}
}
# Set up mock to return our commit data sample
mock_get_recent_commits.return_value = mock_result_commits
# Call the resource endpoint for recent commit status only
result = await get_recent_commits_with_jobs_resource(
"pytorch", "pytorch", "main",
include_success=False,
include_pending=False,
include_failures=False,
page=1,
per_page=10
)
# Verify the resource endpoint was called with the correct arguments
mock_get_recent_commits.assert_called_once_with(
repo_owner="pytorch",
repo_name="pytorch",
branch_or_commit_sha="main",
include_success=False,
include_pending=False,
include_failures=False,
include_commit_details=True,
job_name_filter_regex=None,
failure_line_filter_regex=None,
page=1,
per_page=10,
ctx=None
)
# Result should be a JSON string - parse it back to verify contents
result_data = json.loads(result)
# Check that it contains the expected fields for commit status
self.assertIn("commits", result_data)
self.assertEqual(len(result_data["commits"]), 2)
# The commits shouldn't have job details since we didn't request them
self.assertNotIn("jobs", result_data["commits"][0])
@patch('pytorch_hud.server.mcp_server.get_job_details')
async def test_job_details_resource(self, mock_get_job_details):
"""Test that the job_details_resource properly awaits the async function."""
# Create mock return value
mock_result = {
"job_id": "job123",
"log_url": "https://example.com/logs/job123",
"artifacts": ["artifact1.txt", "artifact2.log"],
"utilization": {"cpu": 80, "memory": 60}
}
# Set up mock to return our sample data
mock_get_job_details.return_value = mock_result
# Call the resource endpoint - use an integer job ID
job_id = 123
result = await get_job_details_resource(job_id)
# Verify the resource endpoint was called with the correct arguments
mock_get_job_details.assert_called_once_with(job_id, ctx=None)
# Result should be a JSON string - parse it back to verify contents
result_data = json.loads(result)
# Check that it contains the expected fields
self.assertEqual(result_data["job_id"], "job123")
self.assertIn("log_url", result_data)
# Removed test_commit_summary_resource, test_job_summary_resource, test_test_summary_resource
# as these functions have been consolidated into get_recent_commits_with_jobs
# Removed test_hud_data_resource as get_hud_data has been removed
if __name__ == "__main__":
unittest.main()