Skip to content

Commit 23f27c1

Browse files
mik-lajmik-laj
authored
Add system tests for CloudKMSHook (#13122)
1 parent cddbf81 commit 23f27c1

File tree

4 files changed

+109
-2
lines changed

4 files changed

+109
-2
lines changed

‎airflow/providers/google/cloud/hooks/kms.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -169,5 +169,4 @@ def decrypt(
169169
metadata=metadata,
170170
)
171171

172-
plaintext = response.plaintext
173-
return plaintext
172+
return response.plaintext

‎tests/providers/google/cloud/hooks/test_kms_system.py‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
import base64
19+
import os
20+
from tempfile import TemporaryDirectory
21+
22+
import pytest
23+
24+
from airflow.providers.google.cloud.hooks.kms import CloudKMSHook
25+
from tests.providers.google.cloud.utils.gcp_authenticator import GCP_KMS_KEY
26+
from tests.test_utils.gcp_system_helpers import GoogleSystemTest, provide_gcp_context
27+
28+
# To prevent resource name collisions, key ring and key resources CANNOT be deleted, so
29+
# to avoid cluttering the project, we only create the key once during project initialization.
30+
# See: https://cloud.google.com/kms/docs/faq#cannot_delete
31+
GCP_KMS_KEYRING_NAME = os.environ.get('GCP_KMS_KEYRING_NAME', 'test-airflow-system-tests-keyring')
32+
GCP_KMS_KEY_NAME = os.environ.get('GCP_KMS_KEY_NAME', 'test-airflow-system-tests-key')
33+
34+
35+
@pytest.mark.credential_file(GCP_KMS_KEY)
36+
class TestKmsHook(GoogleSystemTest):
37+
@provide_gcp_context(GCP_KMS_KEY)
38+
def test_encrypt(self):
39+
with TemporaryDirectory() as tmp_dir:
40+
kms_hook = CloudKMSHook()
41+
content = kms_hook.encrypt(
42+
key_name=(
43+
f"projects/{kms_hook.project_id}/locations/global/keyRings/"
44+
f"{GCP_KMS_KEYRING_NAME}/cryptoKeys/{GCP_KMS_KEY_NAME}"
45+
),
46+
plaintext=b"TEST-SECRET",
47+
)
48+
with open(f"{tmp_dir}/mysecret.txt.encrypted", "wb") as encrypted_file:
49+
encrypted_file.write(base64.b64decode(content))
50+
self.execute_cmd(
51+
[
52+
"gcloud",
53+
"kms",
54+
"decrypt",
55+
"--location",
56+
"global",
57+
"--keyring",
58+
GCP_KMS_KEYRING_NAME,
59+
"--key",
60+
GCP_KMS_KEY_NAME,
61+
"--ciphertext-file",
62+
f"{tmp_dir}/mysecret.txt.encrypted",
63+
"--plaintext-file",
64+
f"{tmp_dir}/mysecret.txt",
65+
]
66+
)
67+
with open(f"{tmp_dir}/mysecret.txt", "rb") as secret_file:
68+
secret = secret_file.read()
69+
self.assertEqual(secret, b"TEST-SECRET")
70+
71+
@provide_gcp_context(GCP_KMS_KEY)
72+
def test_decrypt(self):
73+
with TemporaryDirectory() as tmp_dir:
74+
with open(f"{tmp_dir}/mysecret.txt", "w") as secret_file:
75+
secret_file.write("TEST-SECRET")
76+
self.execute_cmd(
77+
[
78+
"gcloud",
79+
"kms",
80+
"encrypt",
81+
"--location",
82+
"global",
83+
"--keyring",
84+
GCP_KMS_KEYRING_NAME,
85+
"--key",
86+
GCP_KMS_KEY_NAME,
87+
"--plaintext-file",
88+
f"{tmp_dir}/mysecret.txt",
89+
"--ciphertext-file",
90+
f"{tmp_dir}/mysecret.txt.encrypted",
91+
]
92+
)
93+
with open(f"{tmp_dir}/mysecret.txt.encrypted", "rb") as encrypted_file:
94+
encrypted_secret = base64.b64encode(encrypted_file.read()).decode()
95+
96+
kms_hook = CloudKMSHook()
97+
content = kms_hook.decrypt(
98+
key_name=(
99+
f"projects/{kms_hook.project_id}/locations/global/keyRings/"
100+
f"{GCP_KMS_KEYRING_NAME}/cryptoKeys/{GCP_KMS_KEY_NAME}"
101+
),
102+
ciphertext=encrypted_secret,
103+
)
104+
self.assertEqual(content, b"TEST-SECRET")

‎tests/providers/google/cloud/utils/gcp_authenticator.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
GCP_GCS_KEY = 'gcp_gcs.json'
4747
GCP_GCS_TRANSFER_KEY = 'gcp_gcs_transfer.json'
4848
GCP_GKE_KEY = "gcp_gke.json"
49+
GCP_KMS_KEY = "gcp_kms.json"
4950
GCP_LIFE_SCIENCES_KEY = 'gcp_life_sciences.json'
5051
GCP_MEMORYSTORE = 'gcp_memorystore.json'
5152
GCP_PUBSUB_KEY = "gcp_pubsub.json"

‎tests/test_utils/gcp_system_helpers.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,9 +83,12 @@ def provide_gcp_context(
8383
:param scopes: OAuth scopes for the connection
8484
:type scopes: Sequence
8585
:param project_id: The id of GCP project for the connection.
86+
Default: ``os.environ["GCP_PROJECT_ID"]`` or None
8687
:type project_id: str
8788
"""
8889
key_file_path = resolve_full_gcp_key_path(key_file_path) # type: ignore
90+
if project_id is None:
91+
project_id = os.environ.get("GCP_PROJECT_ID")
8992
with provide_gcp_conn_and_credentials(
9093
key_file_path, scopes, project_id
9194
), tempfile.TemporaryDirectory() as gcloud_config_tmp, mock.patch.dict(

0 commit comments

Comments
 (0)