Spaces:
Running
Running
| import asyncio | |
| import unittest | |
| from unittest.mock import AsyncMock, patch | |
| import app | |
| from gradio.mcp import GradioMCPServer | |
| class AdapterTests(unittest.IsolatedAsyncioTestCase): | |
| async def test_only_six_read_tools_exposed(self): | |
| server = GradioMCPServer(app.demo) | |
| self.assertEqual(set(server.tool_to_endpoint), app.ALLOWED) | |
| self.assertEqual(len(app.demo.fns), 6) | |
| for fn in app.demo.fns.values(): | |
| self.assertEqual(fn.concurrency_limit, 4) | |
| self.assertEqual(fn.concurrency_id, 'public_reads') | |
| async def test_unsafe_tool_rejected_before_network(self): | |
| with patch.object(app.httpx, 'AsyncClient') as client: | |
| with self.assertRaisesRegex(ValueError, 'not allowed'): | |
| await app.forward('run', {'dataset': 'lerobot/pusht'}) | |
| client.assert_not_called() | |
| async def test_bounds_and_unknown_parameters_rejected_before_network(self): | |
| for name,args in [('list_datasets',{'limit':1000}),('get_report',{'url':'http://localhost'}), | |
| ('get_findings',{'dataset':'lerobot/pusht'}), | |
| ('search_datasets',{'q':'a'*129}),('search_episodes',{'limit':31}),('search_combined_datasets',{'per_page':1}),('search_combined_datasets',{'direction':'bad'})]: | |
| with self.subTest(name=name),patch.object(app.httpx, 'AsyncClient') as client: | |
| with self.assertRaisesRegex(ValueError, 'Invalid tool arguments'): | |
| await app.forward(name,args) | |
| client.assert_not_called() | |
| async def test_wrappers_preserve_revision_and_pagination(self): | |
| result={'content':[], 'structuredContent':{'report_revision':'r1'},'isError':False} | |
| with patch.object(app,'forward',new=AsyncMock(return_value=result)) as forward: | |
| self.assertEqual(await app.get_findings('r1',dataset='lerobot/pusht',after=5,limit=3),result) | |
| name,args=forward.call_args.args | |
| self.assertEqual(name,'get_findings') | |
| self.assertEqual(args['report_revision'],'r1') | |
| self.assertEqual(args['after'],5) | |
| self.assertEqual(args['limit'],3) | |
| async def test_fixed_anonymous_origin_and_no_environment_auth(self): | |
| with patch.object(app.httpx,'AsyncClient',side_effect=TimeoutError) as client: | |
| with self.assertRaisesRegex(app.gr.Error,'No audit was started'): | |
| await app.get_report(dataset='lerobot/pusht') | |
| self.assertFalse(client.call_args.kwargs['trust_env']) | |
| self.assertFalse(client.call_args.kwargs['follow_redirects']) | |
| self.assertNotIn('headers',client.call_args.kwargs) | |
| self.assertEqual(app.UPSTREAM,'https://mcp.kalanos.ai/mcp') | |
| async def test_transport_preserves_json_and_upstream_errors(self): | |
| import json | |
| from contextlib import asynccontextmanager | |
| from unittest.mock import Mock | |
| async def transport(*args, **kwargs): | |
| yield None, None, None | |
| async def session_context(*args): | |
| yield session | |
| session=Mock(initialize=AsyncMock(),call_tool=AsyncMock()) | |
| for failed in (False,True): | |
| envelope={'content':[], 'structuredContent':{'report_revision':'r1'},'isError':failed} | |
| result=Mock();result.model_dump.return_value=envelope | |
| session.call_tool.return_value=result | |
| with patch.object(app,'streamable_http_client',transport),patch.object(app,'ClientSession',session_context): | |
| if failed: | |
| with self.assertRaisesRegex(app.gr.Error,'Upstream read-only error'): | |
| await app.get_report(report_id='report1') | |
| else: | |
| self.assertEqual(json.loads(await app.get_report(report_id='report1')),envelope) | |
| session.call_tool.assert_awaited_with('get_report',{'report_id':'report1'}) | |
| if __name__ == '__main__': | |
| unittest.main() | |
| class CombinedAdapterTests(unittest.IsolatedAsyncioTestCase): | |
| async def test_combined_wrapper_forwards_filter_and_paging(self): | |
| with patch.object(app, 'forward', new=AsyncMock(return_value='ok')) as f: | |
| self.assertEqual(await app.search_combined_datasets(task=['pick'],robot=['aloha'],page=2,per_page=40,direction='asc',facets=True),'ok') | |
| name,args=f.call_args.args | |
| self.assertEqual(name,'search_combined_datasets') | |
| self.assertEqual([args[k] for k in ['task','robot','page','per_page','direction','facets']],[['pick'],['aloha'],2,40,'asc',True]) | |
| async def test_schema_matches_checked_upstream_contract(self): | |
| import json | |
| from pathlib import Path | |
| from jsonschema import Draft7Validator | |
| schema=app.SCHEMAS['search_combined_datasets'];Draft7Validator.check_schema(schema) | |
| self.assertFalse(list(Draft7Validator(schema).iter_errors({'task':['pick'],'page':2,'per_page':40,'direction':'asc'}))) | |