kaxil commented on code in PR #67217:
URL: https://github.com/apache/airflow/pull/67217#discussion_r3455879118
##########
airflow-core/tests/unit/executors/test_workloads.py:
##########
@@ -171,3 +171,73 @@ def
test_workload_ti_round_trips_through_sdk_generated_model():
assert received.queue == "jdk-17"
assert received.map_index == 3
assert not hasattr(received, "pool_slots")
+
+
+class TestExecuteTaskMakeVersionData:
+ """Tests for ExecuteTask.make() threading version_data through
BundleInfo."""
+
+ @staticmethod
+ def _make_mock_ti(bundle_version, version_data):
+ """Build a mock TI with the attributes ExecuteTask.make() reads."""
+ from unittest.mock import Mock
+
+ ti_id = uuid4()
+ dag_version_id = uuid4()
+
+ ti = Mock()
+ ti.id = ti_id
+ ti.dag_version_id = dag_version_id
+ ti.task_id = "test_task"
+ ti.dag_id = "test_dag"
+ ti.run_id = "test_run"
+ ti.try_number = 1
+ ti.map_index = -1
+ ti.pool_slots = 1
+ ti.queue = "default"
+ ti.priority_weight = 1
+ ti.executor_config = None
+ ti.parent_context_carrier = None
+ ti.context_carrier = None
+ ti.hostname = None
+ ti.external_executor_id = None
+
+ ti.dag_model.bundle_name = "test-bundle"
+ ti.dag_model.relative_fileloc = "dags/test_dag.py"
+
+ ti.dag_run.bundle_version = bundle_version
+
+ if version_data is not None:
+ ti.dag_version.version_data = version_data
+ else:
+ ti.dag_version = None
+
+ return ti
+
+ def test_pinned_run_populates_version_data(self, monkeypatch):
+ """When dag_run.bundle_version is set, version_data from dag_version
flows to BundleInfo."""
+ monkeypatch.setattr(
+ "airflow.utils.helpers.log_filename_template_renderer",
+ lambda: lambda **kwargs: "test.log",
+ )
+
+ version_data = {"schema_version": 1, "files": {"dags/my_dag.py":
"ver123"}}
+ ti = self._make_mock_ti(bundle_version="abc123",
version_data=version_data)
+
+ workload = ExecuteTask.make(ti)
+
+ assert workload.bundle_info.version == "abc123"
+ assert workload.bundle_info.version_data == version_data
+
+ def test_unpinned_run_version_data_is_none(self, monkeypatch):
+ """When dag_run.bundle_version is None (unpinned), version_data must
be None."""
+ monkeypatch.setattr(
+ "airflow.utils.helpers.log_filename_template_renderer",
+ lambda: lambda **kwargs: "test.log",
+ )
+
+ ti = self._make_mock_ti(bundle_version=None, version_data=None)
+
+ workload = ExecuteTask.make(ti)
+
+ assert workload.bundle_info.version is None
+ assert workload.bundle_info.version_data is None
Review Comment:
I think you can just simplify these:
```suggestion
def test_pinned_run_populates_version_data(self, dag_maker, session):
with dag_maker("test_dag", bundle_version="abc123"):
EmptyOperator(task_id="test_task")
dr = dag_maker.create_dagrun()
ti = dr.task_instances[0]
ti.dag_version.version_data = {"schema_version": 1, "files":
{"dags/my_dag.py": "ver123"}}
session.commit()
workload = ExecuteTask.make(ti)
assert workload.bundle_info.version == "abc123"
assert workload.bundle_info.version_data ==
ti.dag_version.version_data
def test_unpinned_run_version_data_is_none(self, dag_maker):
with dag_maker("test_dag", bundle_version=None):
EmptyOperator(task_id="test_task")
dr = dag_maker.create_dagrun()
ti = dr.task_instances[0]
workload = ExecuteTask.make(ti)
assert workload.bundle_info.version is None
assert workload.bundle_info.version_data is None
```
--
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]
