From 8253408afe81c1715a4dec16e509374e21222346 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Mon, 2 Jun 2025 19:03:32 +0530 Subject: [PATCH 1/2] Bring back mapped task extra links test --- .../routes/public/test_extra_links.py | 57 +++++++++++-------- .../tests_common/test_utils/mock_operators.py | 7 +++ 2 files changed, 41 insertions(+), 23 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py index 6bede00104249..848f7fc94cb9b 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py @@ -120,6 +120,9 @@ def _create_dag(self, dag_maker): CustomOperator( task_id=self.task_multiple_links, bash_command=["TEST_LINK_VALUE_1", "TEST_LINK_VALUE_2"] ) + _ = CustomOperator.partial(task_id=self.task_mapped).expand( + bash_command=["TEST_LINK_VALUE_1", "TEST_LINK_VALUE_2"] + ) return dag @pytest.mark.parametrize( @@ -267,29 +270,37 @@ def test_should_respond_200_support_plugins(self, test_client): ).model_dump() ) - @pytest.mark.xfail(reason="TODO: TaskSDK need to fix this, Extra links should work for mapped operator") - def test_should_respond_200_mapped_task_instance(self, test_client): - map_index = 0 - XCom.set( - key="search_query", - value="TEST_LINK_VALUE_1", - task_id=self.task_mapped, - dag_id=self.dag.dag_id, - run_id=self.dag_run_id, - map_index=map_index, - ) - response = test_client.get( - f"/dags/{self.dag_id}/dagRuns/{self.dag_run_id}/taskInstances/{self.task_mapped}/links", - params={"map_index": map_index}, - ) - assert response.status_code == 200 - assert ( - response.json() - == ExtraLinkCollectionResponse( - extra_links={"Google Custom": "http://google.com/custom_base_link?search=TEST_LINK_VALUE_1"}, - total_entries=1, - ).model_dump() - ) + def test_should_respond_200_mapped_task_instance(self, test_client, session): + for map_index, value in enumerate(["TEST_LINK_VALUE_1", "TEST_LINK_VALUE_2"]): + XCom.set( + key="search_query", + value=value, + task_id=self.task_mapped, + dag_id=self.dag_id, + run_id=self.dag_run_id, + map_index=map_index, + ) + XCom.set( + key="_link_CustomOpLink", + value=f"http://google.com/custom_base_link?search={value}", + task_id=self.task_mapped, + dag_id=self.dag_id, + run_id=self.dag_run_id, + map_index=map_index, + ) + session.commit() + response = test_client.get( + f"/dags/{self.dag_id}/dagRuns/{self.dag_run_id}/taskInstances/{self.task_mapped}/links", + params={"map_index": map_index}, + ) + assert response.status_code == 200 + assert ( + response.json() + == ExtraLinkCollectionResponse( + extra_links={"Google Custom": f"http://google.com/custom_base_link?search={value}"}, + total_entries=1, + ).model_dump() + ) def test_should_respond_401_unauthenticated(self, unauthenticated_test_client): response = unauthenticated_test_client.get( diff --git a/devel-common/src/tests_common/test_utils/mock_operators.py b/devel-common/src/tests_common/test_utils/mock_operators.py index 123968e936dec..d4fd43421f6d3 100644 --- a/devel-common/src/tests_common/test_utils/mock_operators.py +++ b/devel-common/src/tests_common/test_utils/mock_operators.py @@ -151,8 +151,15 @@ class CustomOperator(BaseOperator): @property def operator_extra_links(self): """Return operator extra links.""" + # For mapped operators + if not hasattr(self, "bash_command"): + # For mapped operators, we return CustomOpLink since each mapped instance + # will get its own link during runtime + return (CustomOpLink(),) + # For non-mapped operators if isinstance(self.bash_command, str) or self.bash_command is None: return (CustomOpLink(),) + # For operators with multiple commands return (CustomBaseIndexOpLink(i) for i, _ in enumerate(self.bash_command)) def __init__(self, bash_command=None, **kwargs): From 7c6e8feb899e4fec5c81c169be624406c1e75a27 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Fri, 6 Jun 2025 11:45:16 +0530 Subject: [PATCH 2/2] rebase and fixing test --- .../unit/api_fastapi/core_api/routes/public/test_extra_links.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py index fde2263833041..4d84485d48ab4 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py @@ -316,4 +316,4 @@ def test_should_respond_404_invalid_map_index(self, test_client): params={"map_index": 4}, ) assert response.status_code == 404 - assert response.json() == {"detail": "Task with ID = TEST_MAPPED_TASK not found"} + assert response.json() == {"detail": "TaskInstance not found"}