From 3a4739b5bb9c242baa637ae50f61a25e19ebde10 Mon Sep 17 00:00:00 2001
From: phi-friday <phi.friday@gmail.com>
Date: Thu, 13 Jun 2024 10:31:33 +0900
Subject: [PATCH] test: add test test_import_annotations

---
 .../docker/decorators/_with_annotations.py    | 18 +++++++++++++++
 .../docker/decorators/test_docker.py          | 22 +++++++++++++++++++
 2 files changed, 40 insertions(+)
 create mode 100644 tests/providers/docker/decorators/_with_annotations.py

diff --git a/tests/providers/docker/decorators/_with_annotations.py b/tests/providers/docker/decorators/_with_annotations.py
new file mode 100644
index 00000000000000..da68dc4d43f7f0
--- /dev/null
+++ b/tests/providers/docker/decorators/_with_annotations.py
@@ -0,0 +1,18 @@
+from typing import TYPE_CHECKING
+
+from airflow.decorators import task
+
+if TYPE_CHECKING:
+    from airflow.decorators.base import Task
+
+__all__ = []
+
+
+def create_task_factory(image: str, **kwargs: "Any") -> "Task[..., Any]":  # type: ignore  # noqa: F821
+    kwargs.setdefault("multiple_outputs", False)
+
+    @task.docker(image=image, auto_remove="force", **kwargs)
+    def func(value: "dict[str, Any]") -> "Any":  # type: ignore  # noqa: F821
+        assert isinstance(value, dict)
+
+    return func
diff --git a/tests/providers/docker/decorators/test_docker.py b/tests/providers/docker/decorators/test_docker.py
index e4fbe15fc32437..bfd8b57d1837cc 100644
--- a/tests/providers/docker/decorators/test_docker.py
+++ b/tests/providers/docker/decorators/test_docker.py
@@ -284,3 +284,25 @@ def f():
             ret = f()
 
         assert ret.operator.docker_url == "unix://var/run/docker.sock"
+
+    def test_import_annotations(self, dag_maker):
+        from typing import Any
+
+        from airflow.models.dagrun import DagRun  # noqa: TCH001
+        from airflow.utils.state import DagRunState
+
+        from ._with_annotations import create_task_factory  # noqa: TID252
+
+        task_factory = create_task_factory("python:3.9-slim")
+
+        with dag_maker():
+
+            @task.python(multiple_outputs=False)
+            def create_dummy_value() -> dict[str, Any]:
+                return {}
+
+            value = create_dummy_value()
+            _ = task_factory(value)
+
+        dagrun: DagRun = dag_maker.create_dagrun()
+        assert DagRunState(dagrun.state) == DagRunState.SUCCESS