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
@@ -198,34 +197,6 @@ def ydb_server():
198197 stop_container (container )
199198
200199
201- @pytest .fixture (scope = "session" )
202- async def mcp_server (ydb_server ):
203- """Create a YDB MCP server instance for testing."""
204- # Create the server with anonymous credentials
205- server = YDBMCPServer (endpoint = YDB_ENDPOINT , database = YDB_DATABASE )
206-
207- # Store the event loop
208- server ._loop = ensure_event_loop ()
209-
210- try :
211- # Initialize the server by creating the driver
212- await server .create_driver ()
213- yield server
214-
215- # Clean up after tests
216- logger .info ("Cleaning up YDB server resources after tests" )
217- await cleanup_pending_tasks ()
218- if server .driver :
219- await cleanup_driver (server .driver )
220-
221- except Exception as e :
222- logger .error (f"Failed to initialize YDB MCP server: { e } " )
223- pytest .fail (f"Failed to initialize YDB MCP server: { e } " )
224- finally :
225- # Final cleanup
226- await cleanup_pending_tasks ()
227-
228-
229200# Create a global variable to cache the server instance
230201_mcp_server_instance = None
231202
@@ -237,9 +208,7 @@ async def session_mcp_server(ydb_server):
237208
238209 if _mcp_server_instance is None :
239210 # Create the server with anonymous credentials
240- _mcp_server_instance = YDBMCPServer (
241- endpoint = YDB_ENDPOINT , database = YDB_DATABASE , auth_mode = AUTH_MODE_ANONYMOUS
242- )
211+ _mcp_server_instance = YDBMCPServer (endpoint = YDB_ENDPOINT , database = YDB_DATABASE , auth_mode = AUTH_MODE_ANONYMOUS )
243212
244213 try :
245214 # Ensure we have a valid event loop
@@ -278,46 +247,32 @@ async def session_mcp_server(ydb_server):
278247 await cleanup_pending_tasks ()
279248
280249
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
250+ @pytest .fixture (scope = "session" )
251+ async def mcp_server (ydb_server ):
252+ """Create a YDB MCP server instance for testing."""
253+ # Create the server with anonymous credentials
254+ server = YDBMCPServer (endpoint = YDB_ENDPOINT , database = YDB_DATABASE )
287255
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
256+ # Store the event loop
257+ server ._loop = ensure_event_loop ()
292258
293259 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 ()
260+ # Initialize the server by creating the driver
261+ await server .create_driver ()
262+ yield server
301263
302- yield session_mcp_server
264+ # Clean up after tests
265+ logger .info ("Cleaning up YDB server resources after tests" )
266+ await cleanup_pending_tasks ()
267+ if server .driver :
268+ await cleanup_driver (server .driver )
303269
304270 except Exception as e :
305- logger .error (f"Error during test setup : { e } " )
306- pytest .fail (f"Failed to setup test environment : { e } " )
271+ logger .error (f"Failed to initialize YDB MCP server : { e } " )
272+ pytest .fail (f"Failed to initialize YDB MCP server : { e } " )
307273 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 } " )
274+ # Final cleanup
275+ await cleanup_pending_tasks ()
321276
322277
323278async def call_mcp_tool (mcp_server , tool_name , ** params ):
@@ -372,11 +327,7 @@ async def cleanup_after_all_tests():
372327 loop = asyncio .get_running_loop ()
373328 if not loop .is_closed ():
374329 # 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- ]
330+ pending = [task for task in asyncio .all_tasks (loop ) if not task .done () and task != asyncio .current_task ()]
380331
381332 if pending :
382333 logger .debug (f"Cleaning up { len (pending )} pending tasks in final cleanup" )
@@ -394,9 +345,7 @@ async def cleanup_after_all_tests():
394345 for task in pending :
395346 if not task .done ():
396347 with suppress (asyncio .CancelledError , Exception ):
397- task ._log_destroy_pending = (
398- False # Suppress the warning about task destruction
399- )
348+ task ._log_destroy_pending = False # Suppress the warning about task destruction
400349
401350 # Close the loop
402351 loop .stop ()
0 commit comments