This is an automated email from the ASF dual-hosted git repository.
utkarsharma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6b49bb3a4f Fix metadata override for ComputeEngineSSHHook (#37192)
6b49bb3a4f is described below
commit 6b49bb3a4f0994d8beb95b789c07ed894b577abf
Author: M. Olcay Tercanlı <muhammed_terca...@epam.com>
AuthorDate: Wed Feb 7 18:45:05 2024 +0000
Fix metadata override for ComputeEngineSSHHook (#37192)
Co-authored-by: Utkarsh Sharma <utkarshar...@gmail.com>
---
.../providers/google/cloud/hooks/compute_ssh.py | 2 +-
.../google/cloud/hooks/test_compute_ssh.py | 40 ++++++++++++++++++++++
2 files changed, 41 insertions(+), 1 deletion(-)
diff --git a/airflow/providers/google/cloud/hooks/compute_ssh.py
b/airflow/providers/google/cloud/hooks/compute_ssh.py
index 13b1bb3f4b..ee6eb63b02 100644
--- a/airflow/providers/google/cloud/hooks/compute_ssh.py
+++ b/airflow/providers/google/cloud/hooks/compute_ssh.py
@@ -327,7 +327,7 @@ class ComputeEngineSSHHook(SSHHook):
break
else:
new_dict = {"key": "ssh-keys", "value": keys}
- metadata["items"] = [new_dict]
+ metadata["items"] = [*items, new_dict]
self._compute_hook.set_instance_metadata(
zone=self.zone, resource_id=self.instance_name, metadata=metadata,
project_id=self.project_id
diff --git a/tests/providers/google/cloud/hooks/test_compute_ssh.py
b/tests/providers/google/cloud/hooks/test_compute_ssh.py
index ed972afb03..27cfe4fc1b 100644
--- a/tests/providers/google/cloud/hooks/test_compute_ssh.py
+++ b/tests/providers/google/cloud/hooks/test_compute_ssh.py
@@ -507,3 +507,43 @@ class TestComputeEngineHookWithPassedProjectId:
assert isinstance(hook.use_oslogin, bool)
assert 300 == hook.expire_time
assert isinstance(hook.expire_time, int)
+
+ @pytest.mark.parametrize(
+ "metadata, expected_metadata",
+ [
+ ({"items": []}, {"items": [{"key": "ssh-keys", "value":
"user:pubkey\n"}]}),
+ (
+ {"items": [{"key": "test", "value": "test"}]},
+ {"items": [{"key": "ssh-keys", "value": "user:pubkey\n"},
{"key": "test", "value": "test"}]},
+ ),
+ (
+ {"items": [{"key": "ssh-keys", "value": "test"}, {"key":
"test", "value": "test"}]},
+ {
+ "items": [
+ {"key": "ssh-keys", "value": "user:pubkey\ntest"},
+ {"key": "test", "value": "test"},
+ ]
+ },
+ ),
+ ],
+ )
+
@mock.patch("airflow.providers.google.cloud.hooks.compute_ssh.ComputeEngineHook.set_instance_metadata")
+
@mock.patch("airflow.providers.google.cloud.hooks.compute_ssh.ComputeEngineHook.get_instance_info")
+ def test__authorize_compute_engine_instance_metadata(
+ self, mock_get_instance_info, mock_set_instance_metadata, metadata,
expected_metadata
+ ):
+ """Test to ensure the addition metadata is retained"""
+ mock_get_instance_info.return_value = {"metadata": metadata}
+ conn = Connection(
+ conn_type="gcpssh",
+ extra=json.dumps({}),
+ )
+ conn_uri = conn.get_uri()
+ with mock.patch.dict("os.environ", AIRFLOW_CONN_GCPSSH=conn_uri):
+ hook = ComputeEngineSSHHook(gcp_conn_id="gcpssh")
+ hook.user = "user"
+ pubkey = "pubkey"
+ hook._authorize_compute_engine_instance_metadata(pubkey=pubkey)
+
mock_set_instance_metadata.call_args.kwargs["metadata"]["items"].sort(key=lambda
x: x["key"])
+ expected_metadata["items"].sort(key=lambda x: x["key"])
+ assert mock_set_instance_metadata.call_args.kwargs["metadata"] ==
expected_metadata
