harjeevanmaan commented on code in PR #41384:
URL: https://github.com/apache/airflow/pull/41384#discussion_r1747196300
##########
tests/models/test_taskinstance.py:
##########
@@ -5312,3 +5312,65 @@ def
test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_tas
ti.schedule_downstream_tasks()
assert "Error scheduling downstream tasks." in caplog.text
assert "To be swallowed" in caplog.text
+
+
+def test_ti_selector_condition(dag_maker):
+ from airflow.utils.timezone import datetime
+
+ clear_db_runs()
+ files = ["a", "b", "c"]
+
+ start_date = datetime(2024, 1, 1)
+ files = ["file1", "file2", "file3"]
+
+ with dag_maker(dag_id="task_group_mapping_example", start_date=start_date,
schedule=None, catchup=False):
+
+ @task_group(group_id="etl")
+ def etl_pipeline(file):
+ e = EmptyOperator(task_id="e")
+ t = EmptyOperator(task_id="t")
+ last = EmptyOperator(task_id="last")
+
+ e >> t >> last
+
+ etl_pipeline.expand(file=files)
+
+ dag_instance = dag_maker.dag
+ dag_maker.create_dagrun(
+ run_id="manual_run_2024_01_01",
+ state=State.SUCCESS,
+ execution_date=start_date,
+ start_date=start_date,
+ data_interval=(start_date, start_date),
+ run_type=DagRunType.SCHEDULED,
+ )
+
+ # with map_index
+ task_id = "etl.e"
+ task_id_or_regex = [task_id]
+ map_indexes = [0, 1]
+ task_ids = [(task_id, map_index) for map_index in map_indexes]
+ partial_dag = dag_instance.partial_subset(
+ task_ids_or_regex=task_id_or_regex,
+ include_downstream=True,
+ include_upstream=False,
+ )
+
+ # handling downstream tasks
+ if len(partial_dag.task_dict) > 1:
+ task_ids.extend(tid for tid in partial_dag.task_dict if tid != task_id)
+
+ with create_session() as session:
Review Comment:
Already using the `dag_maker `fixture. Added the `session `too
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
