(airflow) branch main updated: Add tests for shift operators in DependencyMixin (#37161)

[husseinawala]

This is an automated email from the ASF dual-hosted git repository.

husseinawala pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new c2cb973bdc Add tests for shift operators in DependencyMixin (#37161)
c2cb973bdc is described below

commit c2cb973bdcad65af342432445d9a33f62e02c30a
Author: raphaelauv <raphael...@gmail.com>
AuthorDate: Thu Feb 8 23:31:06 2024 +0100

    Add tests for shift operators in DependencyMixin (#37161)
    
    * test: taskmin
    
    * clean
    
    * review 1
    
    ---------
    
    Co-authored-by: raphaelauv <raphael...@users.noreply.github.com>
---
 tests/models/test_taskmixin.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tests/models/test_taskmixin.py b/tests/models/test_taskmixin.py
index c34d46fd4e..0e09a05c08 100644
--- a/tests/models/test_taskmixin.py
+++ b/tests/models/test_taskmixin.py
@@ -23,6 +23,7 @@ import pytest
 
 from airflow.decorators import setup, task, teardown
 from airflow.models.baseoperator import BaseOperator
+from airflow.operators.empty import EmptyOperator
 
 pytestmark = pytest.mark.db_test
 
@@ -215,3 +216,31 @@ def 
test_cannot_set_on_failure_fail_dagrun_unless_teardown_taskflow(dag_maker):
             ValueError, match="Cannot mark task 'my_ok_task__2' as setup; task 
is already a teardown."
         ):
             m.operator.is_setup = True
+
+
+class TestDependencyMixin:
+    def test_set_upstream_list(self, dag_maker):
+        with dag_maker("test_set_upstream_list"):
+            op_a = EmptyOperator(task_id="a")
+            op_b = EmptyOperator(task_id="b")
+            op_c = EmptyOperator(task_id="c")
+            op_d = EmptyOperator(task_id="d")
+
+            [op_d, op_c << op_b] << op_a
+
+        assert [op_a] == op_b.upstream_list
+        assert [op_a] == op_d.upstream_list
+        assert [op_b] == op_c.upstream_list
+
+    def test_set_downstream_list(self, dag_maker):
+        with dag_maker("test_set_downstream_list"):
+            op_a = EmptyOperator(task_id="a")
+            op_b = EmptyOperator(task_id="b")
+            op_c = EmptyOperator(task_id="c")
+            op_d = EmptyOperator(task_id="d")
+
+            op_a >> [op_b >> op_c, op_d]
+
+        assert [] == op_b.upstream_list
+        assert [op_a] == op_d.upstream_list
+        assert {op_a, op_b} == set(op_c.upstream_list)