amoghrajesh commented on code in PR #44843:
URL: https://github.com/apache/airflow/pull/44843#discussion_r1883545915
##########
task_sdk/tests/execution_time/test_task_runner.py:
##########
@@ -260,3 +273,98 @@ def test_startup_basic_templated_dag(mocked_parse):
),
log=mock.ANY,
)
+
+
+def test_startup_dag_with_no_templates(mocked_parse):
+ """Test startup of a simple DAG."""
+ from airflow.providers.standard.operators.python import PythonOperator
+
+ task = PythonOperator(
+ task_id="task1",
+ python_callable=lambda: print("hello world!"),
+ )
+
+ what = StartupDetails(
+ ti=TaskInstance(id=uuid7(), task_id="task1", dag_id="basic_dag",
run_id="c", try_number=1),
+ file="",
+ requests_fd=0,
+ )
+ mocked_parse(what, "basic_dag", task)
+
+ with mock.patch(
+ "airflow.sdk.execution_time.task_runner.SUPERVISOR_COMMS", create=True
+ ) as mock_supervisor_comms:
+ mock_supervisor_comms.get_message.return_value = what
+ startup()
+
+ mock_supervisor_comms.send_request.assert_called_once_with(
+ msg=SetRenderedFields(rendered_fields={"op_args": [], "op_kwargs":
{}, "templates_dict": None}),
+ log=mock.ANY,
+ )
+
+
+def test_startup_dag_with_no_templates_mixed_types(mocked_parse):
+ """Test startup of a simple DAG."""
+ from airflow.providers.standard.operators.python import PythonOperator
+
+ task = PythonOperator(
+ task_id="task1",
+ python_callable=lambda *args, **kwargs: print(f"op_args: {args},
op_kwargs: {kwargs}"),
+ op_args=["arg1", "arg2", 1, 2, 3.75, {"key": "value"}],
+ op_kwargs={"key1": "value1", "key2": 99.0, "key3": {"nested_key":
"nested_value"}},
+ )
+
+ what = StartupDetails(
+ ti=TaskInstance(id=uuid7(), task_id="task1", dag_id="basic_dag",
run_id="c", try_number=1),
+ file="",
+ requests_fd=0,
+ )
+ mocked_parse(what, "basic_dag", task)
+
+ with mock.patch(
+ "airflow.sdk.execution_time.task_runner.SUPERVISOR_COMMS", create=True
+ ) as mock_supervisor_comms:
+ mock_supervisor_comms.get_message.return_value = what
+ startup()
+
+ mock_supervisor_comms.send_request.assert_called_once_with(
+ msg=SetRenderedFields(
+ rendered_fields={
+ "templates_dict": None,
+ "op_args": ["arg1", "arg2", 1, 2, 3.75, {"key": "value"}],
+ "op_kwargs": {"key1": "value1", "key2": 99.0, "key3":
{"nested_key": "nested_value"}},
+ }
+ ),
+ log=mock.ANY,
+ )
+
+
+def test_startup_dag_with_tuple_and_set_templated_fields(mocked_parse):
+ """Test startup of a simple DAG with tuple and set templated fields."""
+
+ task = CustomOperator2(
+ task_id="templated_task",
+ my_tup=(
+ 1,
+ 2,
+ ),
+ my_set={1, 2, 3},
+ )
+
+ what = StartupDetails(
+ ti=TaskInstance(id=uuid7(), task_id="templated_task",
dag_id="basic_dag", run_id="c", try_number=1),
+ file="",
+ requests_fd=0,
+ )
+ mocked_parse(what, "basic_dag", task)
+
+ with mock.patch(
+ "airflow.sdk.execution_time.task_runner.SUPERVISOR_COMMS", create=True
+ ) as mock_supervisor_comms:
+ mock_supervisor_comms.get_message.return_value = what
+ startup()
+
+ mock_supervisor_comms.send_request.assert_called_once_with(
+ msg=SetRenderedFields(rendered_fields={"my_set": "{1, 2, 3}",
"my_tup": "(1, 2)"}),
+ log=mock.ANY,
+ )
Review Comment:
Lets not make this change in this PR.
##########
task_sdk/tests/execution_time/test_task_runner.py:
##########
@@ -260,3 +273,98 @@ def test_startup_basic_templated_dag(mocked_parse):
),
log=mock.ANY,
)
+
+
+def test_startup_dag_with_no_templates(mocked_parse):
+ """Test startup of a simple DAG."""
+ from airflow.providers.standard.operators.python import PythonOperator
+
+ task = PythonOperator(
+ task_id="task1",
+ python_callable=lambda: print("hello world!"),
+ )
+
+ what = StartupDetails(
+ ti=TaskInstance(id=uuid7(), task_id="task1", dag_id="basic_dag",
run_id="c", try_number=1),
+ file="",
+ requests_fd=0,
+ )
+ mocked_parse(what, "basic_dag", task)
+
+ with mock.patch(
+ "airflow.sdk.execution_time.task_runner.SUPERVISOR_COMMS", create=True
+ ) as mock_supervisor_comms:
+ mock_supervisor_comms.get_message.return_value = what
+ startup()
+
+ mock_supervisor_comms.send_request.assert_called_once_with(
+ msg=SetRenderedFields(rendered_fields={"op_args": [], "op_kwargs":
{}, "templates_dict": None}),
+ log=mock.ANY,
+ )
+
+
+def test_startup_dag_with_no_templates_mixed_types(mocked_parse):
+ """Test startup of a simple DAG."""
+ from airflow.providers.standard.operators.python import PythonOperator
+
+ task = PythonOperator(
+ task_id="task1",
+ python_callable=lambda *args, **kwargs: print(f"op_args: {args},
op_kwargs: {kwargs}"),
+ op_args=["arg1", "arg2", 1, 2, 3.75, {"key": "value"}],
+ op_kwargs={"key1": "value1", "key2": 99.0, "key3": {"nested_key":
"nested_value"}},
+ )
+
+ what = StartupDetails(
+ ti=TaskInstance(id=uuid7(), task_id="task1", dag_id="basic_dag",
run_id="c", try_number=1),
+ file="",
+ requests_fd=0,
+ )
+ mocked_parse(what, "basic_dag", task)
+
+ with mock.patch(
+ "airflow.sdk.execution_time.task_runner.SUPERVISOR_COMMS", create=True
+ ) as mock_supervisor_comms:
+ mock_supervisor_comms.get_message.return_value = what
+ startup()
+
+ mock_supervisor_comms.send_request.assert_called_once_with(
+ msg=SetRenderedFields(
+ rendered_fields={
+ "templates_dict": None,
+ "op_args": ["arg1", "arg2", 1, 2, 3.75, {"key": "value"}],
+ "op_kwargs": {"key1": "value1", "key2": 99.0, "key3":
{"nested_key": "nested_value"}},
+ }
+ ),
+ log=mock.ANY,
+ )
+
+
+def test_startup_dag_with_tuple_and_set_templated_fields(mocked_parse):
+ """Test startup of a simple DAG with tuple and set templated fields."""
+
+ task = CustomOperator2(
+ task_id="templated_task",
+ my_tup=(
+ 1,
+ 2,
+ ),
+ my_set={1, 2, 3},
+ )
+
+ what = StartupDetails(
+ ti=TaskInstance(id=uuid7(), task_id="templated_task",
dag_id="basic_dag", run_id="c", try_number=1),
+ file="",
+ requests_fd=0,
+ )
+ mocked_parse(what, "basic_dag", task)
+
+ with mock.patch(
+ "airflow.sdk.execution_time.task_runner.SUPERVISOR_COMMS", create=True
+ ) as mock_supervisor_comms:
+ mock_supervisor_comms.get_message.return_value = what
+ startup()
+
+ mock_supervisor_comms.send_request.assert_called_once_with(
+ msg=SetRenderedFields(rendered_fields={"my_set": "{1, 2, 3}",
"my_tup": "(1, 2)"}),
+ log=mock.ANY,
+ )
Review Comment:
Lets not make this change in this PR.
--
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: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
