multiple providers to testing overrides

aioinject/testing/testcontainer.py

@@ -1,6 +1,12 @@
 from __future__ import annotations
 
-from collections.abc import Iterator
+from collections.abc import AsyncIterator, Iterator
+from contextlib import (
+    AsyncExitStack,
+    ExitStack,
+    asynccontextmanager,
+    contextmanager,
+)
 from types import TracebackType
 from typing import TYPE_CHECKING, Any
 
@@ -123,5 +129,22 @@ class TestContainer:
     def __init__(self, container: Container | SyncContainer) -> None:
         self.container = container
 
-    def override(self, provider: Provider[Any]) -> _Override:
-        return _Override(self.container, provider)
+    @asynccontextmanager
+    async def override(self, *providers: Provider[Any]) -> AsyncIterator[None]:
+        overrides = [
+            _Override(self.container, provider) for provider in providers
+        ]
+        async with AsyncExitStack() as exitstack:
+            for override in overrides:
+                await exitstack.enter_async_context(override)
+            yield
+
+    @contextmanager
+    def override_sync(self, *providers: Provider[Any]) -> Iterator[None]:
+        overrides = [
+            _Override(self.container, provider) for provider in providers
+        ]
+        with ExitStack() as exitstack:
+            for override in overrides:
+                exitstack.enter_context(override)
+            yield

tests/testing/test_testcontainer.py

@@ -25,7 +25,7 @@ async def test_override() -> None:
     async with (
         container.context() as ctx,
     ):
-        with testcontainer.override(Object(override)):
+        with testcontainer.override_sync(Object(override)):
             assert await ctx.resolve(int) is override
 
     async with container.context() as ctx:
@@ -153,3 +153,30 @@ class C:
     c = await container.root.resolve(C)
     assert c is not c_mocked
     assert c.b.a is not a_mock
+
+
+async def test_multiple_overrides() -> None:
+    class A:
+        pass
+
+    @dataclasses.dataclass
+    class B:
+        a: A
+
+    @dataclasses.dataclass
+    class C:
+        b: B
+
+    container = Container()
+    container.register(Singleton(A), Singleton(B), Singleton(C))
+    testcontainer = TestContainer(container)
+
+    a_mock = A()
+    b_mock = B(a_mock)
+    async with testcontainer.override(Object(a_mock), Object(b_mock)):
+        c_mocked = await container.root.resolve(C)
+
+    c = await container.root.resolve(C)
+    assert c is not c_mocked
+    assert c.b is not b_mock
+    assert c.b.a is not a_mock