Coverage for src/prisma/testing.py: 100%

1from __future__ import annotations 

2 

3import contextlib 

4from typing import Iterator 

5 

6from . import _registry 

7from .errors import ClientNotRegisteredError 

8from ._registry import RegisteredClient 

9 

10 

11@contextlib.contextmanager 

12def reset_client( 

13 new_client: RegisteredClient | None = None, 

14) -> Iterator[None]: 

15 """Context manager to unregister the current client 

16 

17 Once the context manager exits, the registered client is set back to it's original state 

18 """ 

19 client = _registry._registered_client 

20 if client is None: 

21 raise ClientNotRegisteredError() 

22 

23 try: 

24 _registry._registered_client = new_client 

25 yield 

26 finally: 

27 _registry._registered_client = client 

28 

29 

30def unregister_client() -> None: 

31 """Unregister the current client.""" 

32 if _registry._registered_client is None: 

33 raise ClientNotRegisteredError() 

34 

35 _registry._registered_client = None