Coverage for src/qdrant_loader/cli/asyncio.py: 100%

1""" 

2Re-export asyncio functions to make them mockable in tests. 

3This module provides a simplified interface to asyncio functionality 

4that can be easily mocked in tests. 

5""" 

6 

7import asyncio 

8from asyncio import ( 

9 get_event_loop, 

10 new_event_loop, 

11 run, 

12 set_event_loop, 

13) 

14from collections.abc import Callable, Coroutine 

15from functools import wraps 

16from typing import Any 

17 

18__all__ = ["run", "get_event_loop", "new_event_loop", "set_event_loop"] 

19 

20"""Async utilities for the CLI module.""" 

21 

22 

23def async_command(f: Callable[..., Coroutine[Any, Any, Any]]) -> Callable[..., Any]: 

24 """Decorator to run async commands in the CLI. 

25 

26 This decorator handles the async event loop properly, ensuring that: 

27 1. We don't try to create a new event loop if one already exists 

28 2. We properly clean up the event loop if we create one 

29 3. We handle both testing and production environments correctly 

30 """ 

31 

32 @wraps(f) 

33 def wrapper(*args: Any, **kwargs: Any) -> Any: 

34 try: 

35 loop = asyncio.get_running_loop() 

36 except RuntimeError: 

37 # No running event loop, create a new one 

38 loop = asyncio.new_event_loop() 

39 asyncio.set_event_loop(loop) 

40 should_close = True 

41 else: 

42 should_close = False 

43 

44 try: 

45 return loop.run_until_complete(f(*args, **kwargs)) 

46 finally: 

47 if should_close: 

48 loop.close() 

49 

50 return wrapper