[GitHub] [airflow] casassg commented on a change in pull request #8805: Resolve upstream tasks when template field is XComArg

[GitBox <gi...@apache.org>]

casassg commented on a change in pull request #8805:
URL: https://github.com/apache/airflow/pull/8805#discussion_r431484377



##########
File path: tests/models/test_baseoperator.py
##########
@@ -347,3 +350,66 @@ def test_lineage_composition(self):
         task4 = DummyOperator(task_id="op4", dag=dag)
         task4 > [inlet, outlet, extra]
         self.assertEqual(task4.get_outlet_defs(), [inlet, outlet, extra])
+
+
+@pytest.fixture(scope="class")
+def provide_test_dag_bag():
+    yield DagBag(dag_folder=TEST_DAGS_FOLDER, include_examples=False)
+
+
+class TestXComArgsRelationsAreResolved:
+    def test_upstream_is_set_when_template_field_is_xcomarg(
+        self, provide_test_dag_bag  # pylint: disable=redefined-outer-name
+    ):
+        dag: DAG = provide_test_dag_bag.get_dag("xcomargs_test_1")
+        op1, op2 = sorted(dag.tasks, key=lambda t: t.task_id)
+
+        assert op1 in op2.upstream_list
+        assert op2 in op1.downstream_list
+
+    def test_set_xcomargs_dependencies_works_recursively(
+        self, provide_test_dag_bag  # pylint: disable=redefined-outer-name
+    ):
+        dag: DAG = provide_test_dag_bag.get_dag("xcomargs_test_2")
+        op1, op2, op3, op4 = sorted(dag.tasks, key=lambda t: t.task_id)
+
+        assert op1 in op3.upstream_list
+        assert op2 in op3.upstream_list
+        assert op1 in op4.upstream_list
+        assert op2 in op4.upstream_list
+
+    def test_set_xcomargs_dependencies_works_when_set_after_init(
+        self, provide_test_dag_bag  # pylint: disable=redefined-outer-name
+    ):
+        dag: DAG = provide_test_dag_bag.get_dag("xcomargs_test_3")
+        op1, op2, _ = sorted(dag.tasks, key=lambda t: t.task_id)
+
+        assert op1 in op2.upstream_list
+
+    def test_set_xcomargs_dependencies_no_error_when_outside_dag(self):
+        class CustomOp(DummyOperator):
+            template_fields = ("field",)
+
+            @apply_defaults
+            def __init__(self, field, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.field = field
+
+        op1 = DummyOperator(task_id="op1")
+        CustomOp(task_id="op2", field=op1.output)
+
+    def test_set_xcomargs_dependencies_when_creating_dagbag_with_serialization(self):
+        # Persist DAG
+        dag_id = "xcomargs_test_3"
+        dagbag = DagBag(TEST_DAGS_FOLDER, include_examples=False)

Review comment:
       nit: Why not use `provide_test_dag_bag` here as well?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org