From 24d9ab1ce1152c2862cf5cdbae553c2559055980 Mon Sep 17 00:00:00 2001 From: jx2lee Date: Wed, 29 Jan 2025 14:41:38 +0900 Subject: [PATCH 01/14] template --- airflow/config_templates/config.yml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index 918b111acd43a..da35eb7bf6b76 100644 --- a/airflow/config_templates/config.yml +++ b/airflow/config_templates/config.yml @@ -2595,6 +2595,15 @@ dag_processor: "subdir": "dags", "tracking_ref": "main", "refresh_interval": 0 + }, + { + "name": "my-git-repo", + "classpath": "airflow.dag_processing.bundles.git.GitDagBundle", + "kwargs": { + "subdir": "dags", + "repo_url": "https://github.com/example/my-dags.git", + "tracking_ref": "main", + "refresh_interval": 0 } ] default: > From 07036d3ed2f209069e0cc42cfcbe633e6ae15223 Mon Sep 17 00:00:00 2001 From: jx2lee Date: Wed, 29 Jan 2025 22:51:16 +0900 Subject: [PATCH 02/14] remove repo_url in kwargs --- airflow/config_templates/config.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index da35eb7bf6b76..423eee773cec2 100644 --- a/airflow/config_templates/config.yml +++ b/airflow/config_templates/config.yml @@ -2601,7 +2601,6 @@ dag_processor: "classpath": "airflow.dag_processing.bundles.git.GitDagBundle", "kwargs": { "subdir": "dags", - "repo_url": "https://github.com/example/my-dags.git", "tracking_ref": "main", "refresh_interval": 0 } From 01a35934574f241f9c92441d6eda2ad2bf45ba4f Mon Sep 17 00:00:00 2001 From: jx2lee Date: Thu, 30 Jan 2025 12:19:55 +0900 Subject: [PATCH 03/14] dedup bundle config --- airflow/config_templates/config.yml | 8 -------- 1 file changed, 8 deletions(-) diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index 423eee773cec2..918b111acd43a 100644 --- a/airflow/config_templates/config.yml +++ b/airflow/config_templates/config.yml @@ -2588,14 +2588,6 @@ dag_processor: type: string example: > [ - { - "name": "my-git-repo", - "classpath": "airflow.dag_processing.bundles.git.GitDagBundle", - "kwargs": { - "subdir": "dags", - "tracking_ref": "main", - "refresh_interval": 0 - }, { "name": "my-git-repo", "classpath": "airflow.dag_processing.bundles.git.GitDagBundle", From d4ee6bd98e7981e1fdea314ec79c2e6aa35dee26 Mon Sep 17 00:00:00 2001 From: jx2lee Date: Tue, 28 Jan 2025 23:44:56 +0900 Subject: [PATCH 04/14] inline key --- airflow/dag_processing/bundles/git.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/airflow/dag_processing/bundles/git.py b/airflow/dag_processing/bundles/git.py index abebbb4a33820..dea9b5d81188f 100644 --- a/airflow/dag_processing/bundles/git.py +++ b/airflow/dag_processing/bundles/git.py @@ -19,6 +19,7 @@ import json import os +import tempfile from typing import TYPE_CHECKING, Any from urllib.parse import urlparse @@ -60,6 +61,7 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]: "extra": json.dumps( { "key_file": "optional/path/to/keyfile", + "private_key": "optional/private_key", } ) }, @@ -70,9 +72,15 @@ def __init__(self, git_conn_id="git_default", *args, **kwargs): connection = self.get_connection(git_conn_id) self.repo_url = connection.host self.auth_token = connection.password + self.private_key = connection.extra_dejson.get("private_key") self.key_file = connection.extra_dejson.get("key_file") strict_host_key_checking = connection.extra_dejson.get("strict_host_key_checking", "no") self.env: dict[str, str] = {} + + if self.key_file and self.private_key: + raise AirflowException("Both 'key_file' and 'private_key' cannot be provided at the same time") + if self.private_key: + self._setup_inline_key() if self.key_file: self.env["GIT_SSH_COMMAND"] = ( f"ssh -i {self.key_file} -o IdentitiesOnly=yes -o StrictHostKeyChecking={strict_host_key_checking}" @@ -87,6 +95,13 @@ def _process_git_auth_url(self): elif not self.repo_url.startswith("git@") or not self.repo_url.startswith("https://"): self.repo_url = os.path.expanduser(self.repo_url) + def _setup_inline_key(self): + with tempfile.NamedTemporaryFile(mode="w", delete=False) as tmp_keyfile: + tmp_keyfile.write(self.private_key) + tmp_keyfile.flush() + os.chmod(tmp_keyfile.name, 0o600) + self.tmp_keyfile = tmp_keyfile.name + class GitDagBundle(BaseDagBundle, LoggingMixin): """ From 75d2e1efd54a108d68db8109327e5dd9a155d772 Mon Sep 17 00:00:00 2001 From: jx2lee Date: Tue, 28 Jan 2025 23:52:27 +0900 Subject: [PATCH 05/14] test --- tests/dag_processing/test_dag_bundles.py | 34 ++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/dag_processing/test_dag_bundles.py b/tests/dag_processing/test_dag_bundles.py index 6ab1f3ea68fb6..4a69ff037e5af 100644 --- a/tests/dag_processing/test_dag_bundles.py +++ b/tests/dag_processing/test_dag_bundles.py @@ -108,6 +108,8 @@ def git_repo(tmp_path_factory): CONN_HTTPS = "my_git_conn" CONN_HTTPS_PASSWORD = "my_git_conn_https_password" CONN_ONLY_PATH = "my_git_conn_only_path" +CONN_ONLY_INLINE_KEY = "my_git_conn_only_inline_key" +CONN_BOTH_PATH_INLINE = "my_git_conn_both_path_inline" CONN_NO_REPO_URL = "my_git_conn_no_repo_url" @@ -144,6 +146,28 @@ def setup_class(cls) -> None: conn_id=CONN_ONLY_PATH, host="path/to/repo", conn_type="git", + extra={"key_file": "path/to/key"}, + ) + ) + db.merge_conn( + Connection( + conn_id=CONN_ONLY_INLINE_KEY, + host="path/to/repo", + conn_type="git", + extra={ + "private_key": "inline/key", + }, + ) + ) + db.merge_conn( + Connection( + conn_id=CONN_BOTH_PATH_INLINE, + host="path/to/repo", + conn_type="git", + extra={ + "key_file": "path/to/key", + "private_key": "inline/key", + }, ) ) @@ -179,6 +203,16 @@ def test_env_var(self, session): "GIT_SSH_COMMAND": "ssh -i /files/pkey.pem -o IdentitiesOnly=yes -o StrictHostKeyChecking=yes" } + def test_given_both_private_key_and_key_file(self): + with pytest.raises( + AirflowException, match="Both 'key_file' and 'private_key' cannot be provided at the same time" + ): + GitHook(git_conn_id=CONN_BOTH_PATH_INLINE) + + def test_only_inline_connection_has_tmp_keyfile(self): + hook = GitHook(git_conn_id=CONN_ONLY_INLINE_KEY) + assert hook.tmp_keyfile is not None + class TestGitDagBundle: @classmethod From 1ba36f56e94a9015d4e10f468566e6d4329a997a Mon Sep 17 00:00:00 2001 From: jx2lee Date: Wed, 29 Jan 2025 02:21:42 +0900 Subject: [PATCH 06/14] condition private_key --- airflow/dag_processing/bundles/git.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/dag_processing/bundles/git.py b/airflow/dag_processing/bundles/git.py index dea9b5d81188f..c8fb98e22357a 100644 --- a/airflow/dag_processing/bundles/git.py +++ b/airflow/dag_processing/bundles/git.py @@ -81,7 +81,7 @@ def __init__(self, git_conn_id="git_default", *args, **kwargs): raise AirflowException("Both 'key_file' and 'private_key' cannot be provided at the same time") if self.private_key: self._setup_inline_key() - if self.key_file: + if self.key_file or self.private_key: self.env["GIT_SSH_COMMAND"] = ( f"ssh -i {self.key_file} -o IdentitiesOnly=yes -o StrictHostKeyChecking={strict_host_key_checking}" ) From b4bcb5a5d7db854860689b268736fe25c9c9c387 Mon Sep 17 00:00:00 2001 From: jx2lee Date: Wed, 29 Jan 2025 02:22:44 +0900 Subject: [PATCH 07/14] sample value for private_key --- airflow/dag_processing/bundles/git.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/dag_processing/bundles/git.py b/airflow/dag_processing/bundles/git.py index c8fb98e22357a..9cf3f9f43d0fe 100644 --- a/airflow/dag_processing/bundles/git.py +++ b/airflow/dag_processing/bundles/git.py @@ -61,7 +61,7 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]: "extra": json.dumps( { "key_file": "optional/path/to/keyfile", - "private_key": "optional/private_key", + "private_key": "optional inline private key", } ) }, From 313c1a024e47b0db89fecc3ba107c731c6e01232 Mon Sep 17 00:00:00 2001 From: jx2lee Date: Fri, 31 Jan 2025 23:48:45 +0900 Subject: [PATCH 08/14] refactor --- airflow/dag_processing/bundles/git.py | 47 ++++++++++++++++++--------- 1 file changed, 32 insertions(+), 15 deletions(-) diff --git a/airflow/dag_processing/bundles/git.py b/airflow/dag_processing/bundles/git.py index 9cf3f9f43d0fe..3653520abce1a 100644 --- a/airflow/dag_processing/bundles/git.py +++ b/airflow/dag_processing/bundles/git.py @@ -17,6 +17,7 @@ from __future__ import annotations +import contextlib import json import os import tempfile @@ -74,19 +75,22 @@ def __init__(self, git_conn_id="git_default", *args, **kwargs): self.auth_token = connection.password self.private_key = connection.extra_dejson.get("private_key") self.key_file = connection.extra_dejson.get("key_file") - strict_host_key_checking = connection.extra_dejson.get("strict_host_key_checking", "no") + self.strict_host_key_checking = connection.extra_dejson.get("strict_host_key_checking", "no") self.env: dict[str, str] = {} if self.key_file and self.private_key: raise AirflowException("Both 'key_file' and 'private_key' cannot be provided at the same time") - if self.private_key: - self._setup_inline_key() - if self.key_file or self.private_key: - self.env["GIT_SSH_COMMAND"] = ( - f"ssh -i {self.key_file} -o IdentitiesOnly=yes -o StrictHostKeyChecking={strict_host_key_checking}" - ) + if self.key_file: + self.env["GIT_SSH_COMMAND"] = self._build_ssh_command(self.key_file) self._process_git_auth_url() + def _build_ssh_command(self, key_path: str) -> str: + return ( + f"ssh -i {key_path} " + f"-o IdentitiesOnly=yes " + f"-o StrictHostKeyChecking={self.strict_host_key_checking}" + ) + def _process_git_auth_url(self): if not isinstance(self.repo_url, str): return @@ -95,12 +99,22 @@ def _process_git_auth_url(self): elif not self.repo_url.startswith("git@") or not self.repo_url.startswith("https://"): self.repo_url = os.path.expanduser(self.repo_url) - def _setup_inline_key(self): - with tempfile.NamedTemporaryFile(mode="w", delete=False) as tmp_keyfile: - tmp_keyfile.write(self.private_key) - tmp_keyfile.flush() - os.chmod(tmp_keyfile.name, 0o600) - self.tmp_keyfile = tmp_keyfile.name + def set_git_env(self, key: str) -> dict[str, str]: + if self.key_file: + return self.env + self.env["GIT_SSH_COMMAND"] = self._build_ssh_command(key) + return self.env + + @contextlib.contextmanager + def setup_inline_key(self): + if self.private_key: + with tempfile.NamedTemporaryFile(mode="w", delete=True) as tmp_keyfile: + tmp_keyfile.write(self.private_key) + tmp_keyfile.flush() + os.chmod(tmp_keyfile.name, 0o600) + yield tmp_keyfile.name + else: + yield class GitDagBundle(BaseDagBundle, LoggingMixin): @@ -143,8 +157,11 @@ def __init__( self.log.warning("Could not create GitHook for connection %s : %s", self.git_conn_id, e) def _initialize(self): - self._clone_bare_repo_if_required() - self._ensure_version_in_bare_repo() + with self.hook.setup_inline_key() as tmp_keyfile: + self.hook.env = self.hook.set_git_env(tmp_keyfile) + self._clone_bare_repo_if_required() + self._ensure_version_in_bare_repo() + self._clone_repo_if_required() self.repo.git.checkout(self.tracking_ref) if self.version: From 8fbc641a7a654249bc82bda011eb19d8c03803f2 Mon Sep 17 00:00:00 2001 From: jx2lee Date: Fri, 31 Jan 2025 23:50:38 +0900 Subject: [PATCH 09/14] updated test --- tests/dag_processing/test_dag_bundles.py | 54 ++++++++++++++++++------ 1 file changed, 40 insertions(+), 14 deletions(-) diff --git a/tests/dag_processing/test_dag_bundles.py b/tests/dag_processing/test_dag_bundles.py index 4a69ff037e5af..87f8032e89656 100644 --- a/tests/dag_processing/test_dag_bundles.py +++ b/tests/dag_processing/test_dag_bundles.py @@ -17,6 +17,7 @@ from __future__ import annotations +import os import re import tempfile from pathlib import Path @@ -155,18 +156,7 @@ def setup_class(cls) -> None: host="path/to/repo", conn_type="git", extra={ - "private_key": "inline/key", - }, - ) - ) - db.merge_conn( - Connection( - conn_id=CONN_BOTH_PATH_INLINE, - host="path/to/repo", - conn_type="git", - extra={ - "key_file": "path/to/key", - "private_key": "inline/key", + "private_key": "inline_key", }, ) ) @@ -204,14 +194,50 @@ def test_env_var(self, session): } def test_given_both_private_key_and_key_file(self): + db.merge_conn( + Connection( + conn_id=CONN_BOTH_PATH_INLINE, + host="path/to/repo", + conn_type="git", + extra={ + "key_file": "path/to/key", + "private_key": "inline_key", + }, + ) + ) + with pytest.raises( AirflowException, match="Both 'key_file' and 'private_key' cannot be provided at the same time" ): GitHook(git_conn_id=CONN_BOTH_PATH_INLINE) - def test_only_inline_connection_has_tmp_keyfile(self): + def test_key_file_git_hook_has_env(self): + hook = GitHook(git_conn_id=CONN_ONLY_PATH) + + assert hasattr(hook, "env") + assert hook.env == { + "GIT_SSH_COMMAND": "ssh -i path/to/key -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" + } + + def test_private_key_lazy_env_var(self): hook = GitHook(git_conn_id=CONN_ONLY_INLINE_KEY) - assert hook.tmp_keyfile is not None + assert hook.env == {} + + hook.set_git_env("dummy_inline_key") + assert hook.env == { + "GIT_SSH_COMMAND": "ssh -i dummy_inline_key -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" + } + + def test_setup_inline_key(self): + hook = GitHook(git_conn_id=CONN_ONLY_INLINE_KEY) + assert hasattr(hook, "private_key") + + hook.set_git_env("dummy_inline_key") + + with hook.setup_inline_key() as tmp_keyfile: + assert os.path.exists(tmp_keyfile) + + assert not os.path.exists(tmp_keyfile) class TestGitDagBundle: From 1b775d15fad4b5ba0fa54857121d6275236efb9f Mon Sep 17 00:00:00 2001 From: jx2lee Date: Tue, 4 Feb 2025 20:07:15 +0900 Subject: [PATCH 10/14] adjust refresh --- airflow/dag_processing/bundles/git.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/airflow/dag_processing/bundles/git.py b/airflow/dag_processing/bundles/git.py index 3653520abce1a..40f0516fcf203 100644 --- a/airflow/dag_processing/bundles/git.py +++ b/airflow/dag_processing/bundles/git.py @@ -262,8 +262,11 @@ def _fetch_bare_repo(self): def refresh(self) -> None: if self.version: raise AirflowException("Refreshing a specific version is not supported") - self._fetch_bare_repo() - self.repo.remotes.origin.pull() + + with self.hook.setup_inline_key() as tmp_keyfile: + self.hook.env = self.hook.set_git_env(tmp_keyfile) + self._fetch_bare_repo() + self.repo.remotes.origin.pull() @staticmethod def _convert_git_ssh_url_to_https(url: str) -> str: From c62b86d7efbdce5a32220c007539955491f28edc Mon Sep 17 00:00:00 2001 From: jx2lee Date: Tue, 4 Feb 2025 20:07:53 +0900 Subject: [PATCH 11/14] updated test --- tests/dag_processing/test_dag_bundles.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/dag_processing/test_dag_bundles.py b/tests/dag_processing/test_dag_bundles.py index 87f8032e89656..01309896599bc 100644 --- a/tests/dag_processing/test_dag_bundles.py +++ b/tests/dag_processing/test_dag_bundles.py @@ -147,7 +147,6 @@ def setup_class(cls) -> None: conn_id=CONN_ONLY_PATH, host="path/to/repo", conn_type="git", - extra={"key_file": "path/to/key"}, ) ) db.merge_conn( @@ -212,11 +211,11 @@ def test_given_both_private_key_and_key_file(self): GitHook(git_conn_id=CONN_BOTH_PATH_INLINE) def test_key_file_git_hook_has_env(self): - hook = GitHook(git_conn_id=CONN_ONLY_PATH) + hook = GitHook(git_conn_id=CONN_DEFAULT) assert hasattr(hook, "env") assert hook.env == { - "GIT_SSH_COMMAND": "ssh -i path/to/key -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" + "GIT_SSH_COMMAND": "ssh -i /files/pkey.pem -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" } def test_private_key_lazy_env_var(self): From b0379b84a577b6955c39b9193380cd5006915c3b Mon Sep 17 00:00:00 2001 From: jx2lee Date: Tue, 4 Feb 2025 21:02:57 +0900 Subject: [PATCH 12/14] rename and take env in contextmanager --- airflow/dag_processing/bundles/git.py | 11 ++++++----- tests/dag_processing/test_dag_bundles.py | 10 ++++++---- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/airflow/dag_processing/bundles/git.py b/airflow/dag_processing/bundles/git.py index 40f0516fcf203..bc746e063c200 100644 --- a/airflow/dag_processing/bundles/git.py +++ b/airflow/dag_processing/bundles/git.py @@ -106,14 +106,16 @@ def set_git_env(self, key: str) -> dict[str, str]: return self.env @contextlib.contextmanager - def setup_inline_key(self): + def configure_hook_env(self): if self.private_key: with tempfile.NamedTemporaryFile(mode="w", delete=True) as tmp_keyfile: tmp_keyfile.write(self.private_key) tmp_keyfile.flush() os.chmod(tmp_keyfile.name, 0o600) - yield tmp_keyfile.name + self.set_git_env(tmp_keyfile.name) + yield else: + self.set_git_env(self.private_key) yield @@ -157,7 +159,7 @@ def __init__( self.log.warning("Could not create GitHook for connection %s : %s", self.git_conn_id, e) def _initialize(self): - with self.hook.setup_inline_key() as tmp_keyfile: + with self.hook.configure_hook_env() as tmp_keyfile: self.hook.env = self.hook.set_git_env(tmp_keyfile) self._clone_bare_repo_if_required() self._ensure_version_in_bare_repo() @@ -263,8 +265,7 @@ def refresh(self) -> None: if self.version: raise AirflowException("Refreshing a specific version is not supported") - with self.hook.setup_inline_key() as tmp_keyfile: - self.hook.env = self.hook.set_git_env(tmp_keyfile) + with self.hook.configure_hook_env(): self._fetch_bare_repo() self.repo.remotes.origin.pull() diff --git a/tests/dag_processing/test_dag_bundles.py b/tests/dag_processing/test_dag_bundles.py index 01309896599bc..c4d4ff59ba244 100644 --- a/tests/dag_processing/test_dag_bundles.py +++ b/tests/dag_processing/test_dag_bundles.py @@ -227,16 +227,18 @@ def test_private_key_lazy_env_var(self): "GIT_SSH_COMMAND": "ssh -i dummy_inline_key -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" } - def test_setup_inline_key(self): + def test_configure_hook_env(self): hook = GitHook(git_conn_id=CONN_ONLY_INLINE_KEY) assert hasattr(hook, "private_key") hook.set_git_env("dummy_inline_key") - with hook.setup_inline_key() as tmp_keyfile: - assert os.path.exists(tmp_keyfile) + with hook.configure_hook_env(): + command = hook.env.get("GIT_SSH_COMMAND") + temp_key_path = command.split()[2] + assert os.path.exists(temp_key_path) - assert not os.path.exists(tmp_keyfile) + assert not os.path.exists(temp_key_path) class TestGitDagBundle: From 6b5ed13c68c9fd20b575045a3b795640d01bb34a Mon Sep 17 00:00:00 2001 From: jx2lee Date: Wed, 5 Feb 2025 20:19:44 +0900 Subject: [PATCH 13/14] fixed from comments --- airflow/dag_processing/bundles/git.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/airflow/dag_processing/bundles/git.py b/airflow/dag_processing/bundles/git.py index bc746e063c200..60da6a678ed53 100644 --- a/airflow/dag_processing/bundles/git.py +++ b/airflow/dag_processing/bundles/git.py @@ -80,8 +80,6 @@ def __init__(self, git_conn_id="git_default", *args, **kwargs): if self.key_file and self.private_key: raise AirflowException("Both 'key_file' and 'private_key' cannot be provided at the same time") - if self.key_file: - self.env["GIT_SSH_COMMAND"] = self._build_ssh_command(self.key_file) self._process_git_auth_url() def _build_ssh_command(self, key_path: str) -> str: @@ -99,11 +97,8 @@ def _process_git_auth_url(self): elif not self.repo_url.startswith("git@") or not self.repo_url.startswith("https://"): self.repo_url = os.path.expanduser(self.repo_url) - def set_git_env(self, key: str) -> dict[str, str]: - if self.key_file: - return self.env + def set_git_env(self, key: str) -> None: self.env["GIT_SSH_COMMAND"] = self._build_ssh_command(key) - return self.env @contextlib.contextmanager def configure_hook_env(self): @@ -115,7 +110,7 @@ def configure_hook_env(self): self.set_git_env(tmp_keyfile.name) yield else: - self.set_git_env(self.private_key) + self.set_git_env(self.key_file) yield @@ -159,8 +154,7 @@ def __init__( self.log.warning("Could not create GitHook for connection %s : %s", self.git_conn_id, e) def _initialize(self): - with self.hook.configure_hook_env() as tmp_keyfile: - self.hook.env = self.hook.set_git_env(tmp_keyfile) + with self.hook.configure_hook_env(): self._clone_bare_repo_if_required() self._ensure_version_in_bare_repo() From c7920c989c6bcd3695b0d50333b73937ab11c437 Mon Sep 17 00:00:00 2001 From: jx2lee Date: Wed, 5 Feb 2025 20:21:16 +0900 Subject: [PATCH 14/14] fix test to pass --- tests/dag_processing/test_dag_bundles.py | 29 +++++++++++++----------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/tests/dag_processing/test_dag_bundles.py b/tests/dag_processing/test_dag_bundles.py index c4d4ff59ba244..a9a48139ba9eb 100644 --- a/tests/dag_processing/test_dag_bundles.py +++ b/tests/dag_processing/test_dag_bundles.py @@ -173,11 +173,12 @@ def test_correct_repo_urls(self, conn_id, expected_repo_url): hook = GitHook(git_conn_id=conn_id) assert hook.repo_url == expected_repo_url - def test_env_var(self, session): - hook = GitHook(git_conn_id=CONN_DEFAULT) - assert hook.env == { - "GIT_SSH_COMMAND": "ssh -i /files/pkey.pem -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" - } + def test_env_var_with_configure_hook_env(self, session): + default_hook = GitHook(git_conn_id=CONN_DEFAULT) + with default_hook.configure_hook_env(): + assert default_hook.env == { + "GIT_SSH_COMMAND": "ssh -i /files/pkey.pem -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" + } db.merge_conn( Connection( conn_id="my_git_conn_strict", @@ -187,10 +188,11 @@ def test_env_var(self, session): ) ) - hook = GitHook(git_conn_id="my_git_conn_strict") - assert hook.env == { - "GIT_SSH_COMMAND": "ssh -i /files/pkey.pem -o IdentitiesOnly=yes -o StrictHostKeyChecking=yes" - } + strict_default_hook = GitHook(git_conn_id="my_git_conn_strict") + with strict_default_hook.configure_hook_env(): + assert strict_default_hook.env == { + "GIT_SSH_COMMAND": "ssh -i /files/pkey.pem -o IdentitiesOnly=yes -o StrictHostKeyChecking=yes" + } def test_given_both_private_key_and_key_file(self): db.merge_conn( @@ -210,13 +212,14 @@ def test_given_both_private_key_and_key_file(self): ): GitHook(git_conn_id=CONN_BOTH_PATH_INLINE) - def test_key_file_git_hook_has_env(self): + def test_key_file_git_hook_has_env_with_configure_hook_env(self): hook = GitHook(git_conn_id=CONN_DEFAULT) assert hasattr(hook, "env") - assert hook.env == { - "GIT_SSH_COMMAND": "ssh -i /files/pkey.pem -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" - } + with hook.configure_hook_env(): + assert hook.env == { + "GIT_SSH_COMMAND": "ssh -i /files/pkey.pem -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" + } def test_private_key_lazy_env_var(self): hook = GitHook(git_conn_id=CONN_ONLY_INLINE_KEY)