1515from urllib .parse import urlparse
1616
1717import pytest
18- import ydb
1918
2019from tests .docker_utils import start_ydb_container , stop_container , wait_for_port
2120from ydb_mcp .server import AUTH_MODE_ANONYMOUS , YDBMCPServer
@@ -237,9 +236,7 @@ async def session_mcp_server(ydb_server):
237236
238237 if _mcp_server_instance is None :
239238 # Create the server with anonymous credentials
240- _mcp_server_instance = YDBMCPServer (
241- endpoint = YDB_ENDPOINT , database = YDB_DATABASE , auth_mode = AUTH_MODE_ANONYMOUS
242- )
239+ _mcp_server_instance = YDBMCPServer (endpoint = YDB_ENDPOINT , database = YDB_DATABASE , auth_mode = AUTH_MODE_ANONYMOUS )
243240
244241 try :
245242 # Ensure we have a valid event loop
@@ -278,46 +275,32 @@ async def session_mcp_server(ydb_server):
278275 await cleanup_pending_tasks ()
279276
280277
281- @pytest .fixture (scope = "function" )
282- async def mcp_server (session_mcp_server ):
283- """Provide a clean MCP server connection for each test by restarting the connection."""
284- if session_mcp_server is None :
285- pytest .fail ("Could not get a valid MCP server instance" )
286- return
278+ @pytest .fixture (scope = "session" )
279+ async def mcp_server (ydb_server ): # noqa: F811
280+ """Create a YDB MCP server instance for testing."""
281+ # Create the server with anonymous credentials
282+ server = YDBMCPServer (endpoint = YDB_ENDPOINT , database = YDB_DATABASE )
287283
288- # Reset server state to default
289- session_mcp_server .auth_mode = AUTH_MODE_ANONYMOUS
290- session_mcp_server .login = None
291- session_mcp_server .password = None
284+ # Store the event loop
285+ server ._loop = ensure_event_loop ()
292286
293287 try :
294- # Clean up any leftover tasks before restart
295- await cleanup_pending_tasks ()
296-
297- # Restart the connection to ensure clean environment for the test
298- if session_mcp_server .driver is not None :
299- logger .info ("Restarting YDB connection for clean test environment" )
300- await session_mcp_server .restart ()
288+ # Initialize the server by creating the driver
289+ await server .create_driver ()
290+ yield server
301291
302- yield session_mcp_server
292+ # Clean up after tests
293+ logger .info ("Cleaning up YDB server resources after tests" )
294+ await cleanup_pending_tasks ()
295+ if server .driver :
296+ await cleanup_driver (server .driver )
303297
304298 except Exception as e :
305- logger .error (f"Error during test setup : { e } " )
306- pytest .fail (f"Failed to setup test environment : { e } " )
299+ logger .error (f"Failed to initialize YDB MCP server : { e } " )
300+ pytest .fail (f"Failed to initialize YDB MCP server : { e } " )
307301 finally :
308- # Reset server state after test
309- try :
310- session_mcp_server .auth_mode = AUTH_MODE_ANONYMOUS
311- session_mcp_server .login = None
312- session_mcp_server .password = None
313-
314- # Clean up any tasks from the test
315- await cleanup_pending_tasks ()
316-
317- # Restart to clean state
318- await session_mcp_server .restart ()
319- except Exception as e :
320- logger .error (f"Error during test cleanup: { e } " )
302+ # Final cleanup
303+ await cleanup_pending_tasks ()
321304
322305
323306async def call_mcp_tool (mcp_server , tool_name , ** params ):
@@ -372,11 +355,7 @@ async def cleanup_after_all_tests():
372355 loop = asyncio .get_running_loop ()
373356 if not loop .is_closed ():
374357 # Cancel all tasks
375- pending = [
376- task
377- for task in asyncio .all_tasks (loop )
378- if not task .done () and task != asyncio .current_task ()
379- ]
358+ pending = [task for task in asyncio .all_tasks (loop ) if not task .done () and task != asyncio .current_task ()]
380359
381360 if pending :
382361 logger .debug (f"Cleaning up { len (pending )} pending tasks in final cleanup" )
@@ -394,9 +373,7 @@ async def cleanup_after_all_tests():
394373 for task in pending :
395374 if not task .done ():
396375 with suppress (asyncio .CancelledError , Exception ):
397- task ._log_destroy_pending = (
398- False # Suppress the warning about task destruction
399- )
376+ task ._log_destroy_pending = False # Suppress the warning about task destruction
400377
401378 # Close the loop
402379 loop .stop ()
0 commit comments