diff --git a/providers/google/src/airflow/providers/google/cloud/links/base.py b/providers/google/src/airflow/providers/google/cloud/links/base.py index ea6d8850c95f5..6a46f3946d7d7 100644 --- a/providers/google/src/airflow/providers/google/cloud/links/base.py +++ b/providers/google/src/airflow/providers/google/cloud/links/base.py @@ -89,7 +89,12 @@ def persist(cls, context: Context, **value): def get_config(self, operator, ti_key): conf = {} conf.update(getattr(operator, "extra_links_params", {})) - conf.update(XCom.get_value(key=self.key, ti_key=ti_key) or {}) + + xcom_value = XCom.get_value(key=self.key, ti_key=ti_key) or {} + if isinstance(xcom_value, str): + xcom_value = {self.key: xcom_value} + + conf.update(xcom_value) # if the config did not define, return None to stop URL formatting if not conf: diff --git a/providers/google/tests/unit/google/cloud/links/test_base_link.py b/providers/google/tests/unit/google/cloud/links/test_base_link.py index b74963d916ed3..3e369c1b86323 100644 --- a/providers/google/tests/unit/google/cloud/links/test_base_link.py +++ b/providers/google/tests/unit/google/cloud/links/test_base_link.py @@ -128,3 +128,35 @@ def test_get_link(self, create_task_instance_of_operator, session, mock_supervis ) actual_url = link.get_link(operator=ti.task, ti_key=ti.key) assert actual_url == expected_url + + def test_get_config_with_string_xcom_value(self): + link = GoogleLink() + + mock_operator = mock.MagicMock() + mock_operator.extra_links_params = {} + + mock_ti_key = mock.MagicMock() + + test_string_value = "some_string_value" + + with mock.patch("airflow.providers.google.cloud.links.base.XCom.get_value") as mock_get_value: + # Test case 1: XCom returns a string + mock_get_value.return_value = test_string_value + config = link.get_config(mock_operator, mock_ti_key) + + expected_config = {link.key: test_string_value, "namespace": "default"} + assert config == expected_config + + # Test case 2: XCom returns a dictionary (normal case) + test_dict_value = {"key1": "value1", "key2": "value2"} + mock_get_value.return_value = test_dict_value + config = link.get_config(mock_operator, mock_ti_key) + + expected_config = {**test_dict_value, "namespace": "default"} + assert config == expected_config + + # Test case 3: XCom returns None + mock_get_value.return_value = None + config = link.get_config(mock_operator, mock_ti_key) + + assert config is None