From 41777c7c4300b0591a1e0cb5ce7f3f438136bd20 Mon Sep 17 00:00:00 2001
From: Artur Shiriev <lesnik512@gmail.com>
Date: Sat, 17 Oct 2026 20:22:56 +0300
Subject: [PATCH] fix: resolve a task's dependencies from its own app's
 container

---
 modern_di_celery/main.py | 11 +++++++----
 tests/test_inject.py     | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/modern_di_celery/main.py b/modern_di_celery/main.py
index 72827bc..b4a4a1c 100644
--- a/modern_di_celery/main.py
+++ b/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()
diff --git a/tests/test_inject.py b/tests/test_inject.py
index d13dd1d..430a995 100644
--- a/tests/test_inject.py
+++ b/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"