From 86205cc95f9002283ec8902eb2749d47afb3cf57 Mon Sep 17 00:00:00 2001 From: Danny Mccormick Date: Mon, 1 Dec 2025 15:58:12 -0500 Subject: [PATCH 1/4] Reuse secret instead of creating every time --- .../apache_beam/transforms/util_test.py | 127 +++++++++++++++++- 1 file changed, 120 insertions(+), 7 deletions(-) diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py index 34e251fad1c7..e590de0f0115 100644 --- a/sdks/python/apache_beam/transforms/util_test.py +++ b/sdks/python/apache_beam/transforms/util_test.py @@ -72,6 +72,7 @@ from apache_beam.transforms.core import FlatMapTuple from apache_beam.transforms.trigger import AfterCount from apache_beam.transforms.trigger import Repeatedly +from apache_beam.transforms.util import GcpHsmGeneratedSecret from apache_beam.transforms.util import GcpSecret from apache_beam.transforms.util import Secret from apache_beam.transforms.window import FixedWindows @@ -321,8 +322,7 @@ def setUpClass(cls): if secretmanager is not None: cls.project_id = 'apache-beam-testing' py_version = f'_py{sys.version_info.major}{sys.version_info.minor}' - secret_postfix = datetime.now().strftime('%m%d_%H%M%S') + py_version - cls.secret_id = 'gbek_util_secret_tests_' + secret_postfix + cls.secret_id = 'gbek_util_secret_tests' cls.client = secretmanager.SecretManagerServiceClient() cls.project_path = f'projects/{cls.project_id}' cls.secret_path = f'{cls.project_path}/secrets/{cls.secret_id}' @@ -350,11 +350,6 @@ def setUpClass(cls): cls.gcp_secret = GcpSecret(version_name) cls.secret_option = f'type:GcpSecret;version_name:{version_name}' - @classmethod - def tearDownClass(cls): - if secretmanager is not None: - cls.client.delete_secret(request={'name': cls.secret_path}) - def test_gbek_fake_secret_manager_roundtrips(self): fakeSecret = FakeSecret() @@ -439,6 +434,124 @@ def test_gbek_gcp_secret_manager_throws(self): result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) +@unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') +class GcpHsmGeneratedSecretTest(unittest.TestCase): + def setUp(self): + self.mock_secret_manager_client = mock.MagicMock() + self.mock_kms_client = mock.MagicMock() + + # Patch the clients + self.secretmanager_patcher = mock.patch( + 'google.cloud.secretmanager.SecretManagerServiceClient', + return_value=self.mock_secret_manager_client) + self.kms_patcher = mock.patch( + 'google.cloud.kms.KeyManagementServiceClient', + return_value=self.mock_kms_client) + self.os_urandom_patcher = mock.patch('os.urandom', return_value=b'0' * 32) + self.hkdf_patcher = mock.patch( + 'cryptography.hazmat.primitives.kdf.hkdf.HKDF.derive', + return_value=b'derived_key') + + self.secretmanager_patcher.start() + self.kms_patcher.start() + self.os_urandom_patcher.start() + self.hkdf_patcher.start() + + def tearDown(self): + self.secretmanager_patcher.stop() + self.kms_patcher.stop() + self.os_urandom_patcher.stop() + self.hkdf_patcher.stop() + + def test_happy_path_secret_creation(self): + from google.api_core import exceptions as api_exceptions + + project_id = 'test-project' + location_id = 'global' + key_ring_id = 'test-key-ring' + key_id = 'test-key' + job_name = 'test-job' + + secret = GcpHsmGeneratedSecret( + project_id, location_id, key_ring_id, key_id, job_name) + + # Mock responses for secret creation path + self.mock_secret_manager_client.access_secret_version.side_effect = [ + api_exceptions.NotFound('not found'), # first check + api_exceptions.NotFound('not found'), # second check + mock.MagicMock(payload=mock.MagicMock(data=b'derived_key')) + ] + self.mock_kms_client.encrypt.return_value = mock.MagicMock( + ciphertext=b'encrypted_nonce') + + secret_bytes = secret.get_secret_bytes() + self.assertEqual(secret_bytes, b'derived_key') + + # Assertions on mocks + secret_version_path = ( + f'projects/{project_id}/secrets/{secret._secret_version_name}' + '/versions/1') + self.mock_secret_manager_client.access_secret_version.assert_any_call( + request={'name': secret_version_path}) + self.assertEqual( + self.mock_secret_manager_client.access_secret_version.call_count, 3) + self.mock_secret_manager_client.create_secret.assert_called_once() + self.mock_kms_client.encrypt.assert_called_once() + self.mock_secret_manager_client.add_secret_version.assert_called_once() + + def test_secret_already_exists(self): + from google.api_core import exceptions as api_exceptions + + project_id = 'test-project' + location_id = 'global' + key_ring_id = 'test-key-ring' + key_id = 'test-key' + job_name = 'test-job' + + secret = GcpHsmGeneratedSecret( + project_id, location_id, key_ring_id, key_id, job_name) + + # Mock responses for secret creation path + self.mock_secret_manager_client.access_secret_version.side_effect = [ + api_exceptions.NotFound('not found'), + api_exceptions.NotFound('not found'), + mock.MagicMock(payload=mock.MagicMock(data=b'derived_key')) + ] + self.mock_secret_manager_client.create_secret.side_effect = ( + api_exceptions.AlreadyExists('exists')) + self.mock_kms_client.encrypt.return_value = mock.MagicMock( + ciphertext=b'encrypted_nonce') + + secret_bytes = secret.get_secret_bytes() + self.assertEqual(secret_bytes, b'derived_key') + + # Assertions on mocks + self.mock_secret_manager_client.create_secret.assert_called_once() + self.mock_secret_manager_client.add_secret_version.assert_called_once() + + def test_secret_version_already_exists(self): + project_id = 'test-project' + location_id = 'global' + key_ring_id = 'test-key-ring' + key_id = 'test-key' + job_name = 'test-job' + + secret = GcpHsmGeneratedSecret( + project_id, location_id, key_ring_id, key_id, job_name) + + self.mock_secret_manager_client.access_secret_version.return_value = ( + mock.MagicMock(payload=mock.MagicMock(data=b'existing_dek'))) + + secret_bytes = secret.get_secret_bytes() + self.assertEqual(secret_bytes, b'existing_dek') + + # Assertions + self.mock_secret_manager_client.access_secret_version.assert_called_once() + self.mock_secret_manager_client.create_secret.assert_not_called() + self.mock_secret_manager_client.add_secret_version.assert_not_called() + self.mock_kms_client.encrypt.assert_not_called() + + class FakeClock(object): def __init__(self, now=time.time()): self._now = now From d7baa074165f146dd334f22a5a256a4cd3a18f15 Mon Sep 17 00:00:00 2001 From: Danny Mccormick Date: Mon, 1 Dec 2025 16:20:03 -0500 Subject: [PATCH 2/4] Reuse secret instead of creating every time --- .../apache_beam/transforms/util_test.py | 119 ------------------ 1 file changed, 119 deletions(-) diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py index e590de0f0115..e15594aba8ba 100644 --- a/sdks/python/apache_beam/transforms/util_test.py +++ b/sdks/python/apache_beam/transforms/util_test.py @@ -72,7 +72,6 @@ from apache_beam.transforms.core import FlatMapTuple from apache_beam.transforms.trigger import AfterCount from apache_beam.transforms.trigger import Repeatedly -from apache_beam.transforms.util import GcpHsmGeneratedSecret from apache_beam.transforms.util import GcpSecret from apache_beam.transforms.util import Secret from apache_beam.transforms.window import FixedWindows @@ -434,124 +433,6 @@ def test_gbek_gcp_secret_manager_throws(self): result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) -@unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') -class GcpHsmGeneratedSecretTest(unittest.TestCase): - def setUp(self): - self.mock_secret_manager_client = mock.MagicMock() - self.mock_kms_client = mock.MagicMock() - - # Patch the clients - self.secretmanager_patcher = mock.patch( - 'google.cloud.secretmanager.SecretManagerServiceClient', - return_value=self.mock_secret_manager_client) - self.kms_patcher = mock.patch( - 'google.cloud.kms.KeyManagementServiceClient', - return_value=self.mock_kms_client) - self.os_urandom_patcher = mock.patch('os.urandom', return_value=b'0' * 32) - self.hkdf_patcher = mock.patch( - 'cryptography.hazmat.primitives.kdf.hkdf.HKDF.derive', - return_value=b'derived_key') - - self.secretmanager_patcher.start() - self.kms_patcher.start() - self.os_urandom_patcher.start() - self.hkdf_patcher.start() - - def tearDown(self): - self.secretmanager_patcher.stop() - self.kms_patcher.stop() - self.os_urandom_patcher.stop() - self.hkdf_patcher.stop() - - def test_happy_path_secret_creation(self): - from google.api_core import exceptions as api_exceptions - - project_id = 'test-project' - location_id = 'global' - key_ring_id = 'test-key-ring' - key_id = 'test-key' - job_name = 'test-job' - - secret = GcpHsmGeneratedSecret( - project_id, location_id, key_ring_id, key_id, job_name) - - # Mock responses for secret creation path - self.mock_secret_manager_client.access_secret_version.side_effect = [ - api_exceptions.NotFound('not found'), # first check - api_exceptions.NotFound('not found'), # second check - mock.MagicMock(payload=mock.MagicMock(data=b'derived_key')) - ] - self.mock_kms_client.encrypt.return_value = mock.MagicMock( - ciphertext=b'encrypted_nonce') - - secret_bytes = secret.get_secret_bytes() - self.assertEqual(secret_bytes, b'derived_key') - - # Assertions on mocks - secret_version_path = ( - f'projects/{project_id}/secrets/{secret._secret_version_name}' - '/versions/1') - self.mock_secret_manager_client.access_secret_version.assert_any_call( - request={'name': secret_version_path}) - self.assertEqual( - self.mock_secret_manager_client.access_secret_version.call_count, 3) - self.mock_secret_manager_client.create_secret.assert_called_once() - self.mock_kms_client.encrypt.assert_called_once() - self.mock_secret_manager_client.add_secret_version.assert_called_once() - - def test_secret_already_exists(self): - from google.api_core import exceptions as api_exceptions - - project_id = 'test-project' - location_id = 'global' - key_ring_id = 'test-key-ring' - key_id = 'test-key' - job_name = 'test-job' - - secret = GcpHsmGeneratedSecret( - project_id, location_id, key_ring_id, key_id, job_name) - - # Mock responses for secret creation path - self.mock_secret_manager_client.access_secret_version.side_effect = [ - api_exceptions.NotFound('not found'), - api_exceptions.NotFound('not found'), - mock.MagicMock(payload=mock.MagicMock(data=b'derived_key')) - ] - self.mock_secret_manager_client.create_secret.side_effect = ( - api_exceptions.AlreadyExists('exists')) - self.mock_kms_client.encrypt.return_value = mock.MagicMock( - ciphertext=b'encrypted_nonce') - - secret_bytes = secret.get_secret_bytes() - self.assertEqual(secret_bytes, b'derived_key') - - # Assertions on mocks - self.mock_secret_manager_client.create_secret.assert_called_once() - self.mock_secret_manager_client.add_secret_version.assert_called_once() - - def test_secret_version_already_exists(self): - project_id = 'test-project' - location_id = 'global' - key_ring_id = 'test-key-ring' - key_id = 'test-key' - job_name = 'test-job' - - secret = GcpHsmGeneratedSecret( - project_id, location_id, key_ring_id, key_id, job_name) - - self.mock_secret_manager_client.access_secret_version.return_value = ( - mock.MagicMock(payload=mock.MagicMock(data=b'existing_dek'))) - - secret_bytes = secret.get_secret_bytes() - self.assertEqual(secret_bytes, b'existing_dek') - - # Assertions - self.mock_secret_manager_client.access_secret_version.assert_called_once() - self.mock_secret_manager_client.create_secret.assert_not_called() - self.mock_secret_manager_client.add_secret_version.assert_not_called() - self.mock_kms_client.encrypt.assert_not_called() - - class FakeClock(object): def __init__(self, now=time.time()): self._now = now From 8518c5dd779c7acd7fc565b16cee6a91ac6ff784 Mon Sep 17 00:00:00 2001 From: Danny Mccormick Date: Mon, 1 Dec 2025 16:59:51 -0500 Subject: [PATCH 3/4] lint --- sdks/python/apache_beam/transforms/util_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py index e15594aba8ba..61d4a806f140 100644 --- a/sdks/python/apache_beam/transforms/util_test.py +++ b/sdks/python/apache_beam/transforms/util_test.py @@ -320,7 +320,6 @@ class GroupByEncryptedKeyTest(unittest.TestCase): def setUpClass(cls): if secretmanager is not None: cls.project_id = 'apache-beam-testing' - py_version = f'_py{sys.version_info.major}{sys.version_info.minor}' cls.secret_id = 'gbek_util_secret_tests' cls.client = secretmanager.SecretManagerServiceClient() cls.project_path = f'projects/{cls.project_id}' From 5b4ef05ac1ac6dbcaf5177dfbe0e70da6ade2c36 Mon Sep 17 00:00:00 2001 From: Danny Mccormick Date: Mon, 1 Dec 2025 21:20:20 -0500 Subject: [PATCH 4/4] lint --- sdks/python/apache_beam/transforms/util_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py index 61d4a806f140..3dc089134a39 100644 --- a/sdks/python/apache_beam/transforms/util_test.py +++ b/sdks/python/apache_beam/transforms/util_test.py @@ -28,7 +28,6 @@ import math import random import re -import sys import time import unittest import warnings