fix: resolve a task's dependencies from its own app's container

modern_di_celery/main.py

@@ -2,7 +2,7 @@
 import typing
 import weakref
 
-from celery import Celery, Task, current_app, signals
+from celery import Celery, Task, current_app, current_task, signals
 from celery.local import Proxy
 from celery.utils.dispatch import Signal
 from celery.utils.functional import head_from_fun
@@ -22,6 +22,11 @@ def _current_app() -> Celery:
     return _real_app(typing.cast(Celery, current_app))
 
 
+def _task_app() -> Celery:
+    task = current_task._get_current_object()  # noqa: SLF001
+    return _current_app() if task is None else _real_app(task.app)
+
+
 def _worker_app(sender: object) -> Celery:
     app = getattr(sender, "app", None)
     return _current_app() if app is None else _real_app(app)
@@ -125,9 +130,7 @@ def inject(func: typing.Callable[..., T]) -> typing.Callable[..., T]:
     visible_signature = signature.replace(parameters=visible_params)
 
     def wrapper(*args: typing.Any, **kwargs: typing.Any) -> T:  # noqa: ANN401
-        with fetch_di_container(typing.cast(Celery, current_app)).build_child_container(
-            scope=Scope.REQUEST
-        ) as container:
+        with fetch_di_container(_task_app()).build_child_container(scope=Scope.REQUEST) as container:
             resolved = integrations.resolve_markers(container, di_params)
             bound = visible_signature.bind(*args, **kwargs)
             bound.apply_defaults()

tests/test_inject.py

@@ -5,7 +5,7 @@
 from celery import Celery
 from modern_di import Container, Group, Scope, providers
 
-from modern_di_celery import FromDI, inject, setup_di
+from modern_di_celery import DITask, FromDI, inject, setup_di
 from tests.dependencies import Dependencies, DependentCreator, SimpleCreator
 
 
@@ -138,3 +138,38 @@ def sample(_svc: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)]) -> None
 
     with pytest.raises(RuntimeError, match=r"setup_di\(app, container\)"):
         sample.delay().get()
+
+
+@pytest.mark.parametrize("wiring", ["inject", "DITask"])
+def test_a_task_resolves_from_its_own_app_when_another_app_is_current(wiring: str) -> None:
+    class Names(Group):
+        name = providers.Factory(creator=SimpleCreator, kwargs={"dep1": "unset"})
+
+    def make_app(main: str) -> Celery:
+        app_ = Celery(main)
+        app_.conf.task_always_eager = True
+        container = Container(groups=[Names])
+        container.override(Names.name, SimpleCreator(dep1=main))
+        setup_di(app_, container)
+        return app_
+
+    owner = make_app("owner")
+    make_app("other").set_current()
+
+    def read_name(name: typing.Annotated[SimpleCreator, FromDI(Names.name)]) -> str:
+        return name.dep1
+
+    task = owner.task(inject(read_name)) if wiring == "inject" else owner.task(base=DITask)(read_name)
+
+    assert task.delay().get() == "owner"
+    assert task() == "owner"
+
+
+def test_inject_outside_a_task_resolves_from_the_current_app(app: Celery) -> None:
+    @inject
+    def read_dep1(instance: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)]) -> str:
+        return instance.dep1
+
+    app.set_current()
+
+    assert read_dep1() == "original"