def test_get_conn_uri(self, connections_prefix, mock_client_callable,
mock_get_creds):
mock_get_creds.return_value = CREDENTIALS, PROJECT_ID
mock_client = mock.MagicMock()
mock_client_callable.return_value = mock_client
test_response = AccessSecretVersionResponse()
test_response.payload.data = CONN_URI.encode("UTF-8")
mock_client.access_secret_version.return_value = test_response
secrets_manager_backend = CloudSecretsManagerBackend(
connections_prefix=connections_prefix)
returned_uri = secrets_manager_backend.get_conn_uri(conn_id=CONN_ID)
self.assertEqual(CONN_URI, returned_uri)
mock_client.secret_version_path.assert_called_once_with(
PROJECT_ID, f"{connections_prefix}/{CONN_ID}", "latest")
def test_get_variable(self, variables_prefix, mock_client_callable,
mock_get_creds):
mock_get_creds.return_value = CREDENTIALS, PROJECT_ID
mock_client = mock.MagicMock()
mock_client_callable.return_value = mock_client
test_response = AccessSecretVersionResponse()
test_response.payload.data = VAR_VALUE.encode("UTF-8")
mock_client.access_secret_version.return_value = test_response
secrets_manager_backend = CloudSecretsManagerBackend(
variables_prefix=variables_prefix)
secret_id = secrets_manager_backend.build_path(variables_prefix,
VAR_KEY, SEP)
returned_uri = secrets_manager_backend.get_variable(VAR_KEY)
self.assertEqual(VAR_VALUE, returned_uri)
mock_client.secret_version_path.assert_called_once_with(
PROJECT_ID, secret_id, "latest")
def test_get_variable_non_existent_key(self, mock_client_callable,
mock_get_creds):
mock_get_creds.return_value = CREDENTIALS, PROJECT_ID
mock_client = mock.MagicMock()
mock_client_callable.return_value = mock_client
# The requested secret id or secret version does not exist
mock_client.access_secret_version.side_effect = NotFound('test-msg')
secrets_manager_backend = CloudSecretsManagerBackend(
variables_prefix=VARIABLES_PREFIX)
secret_id = secrets_manager_backend.build_path(VARIABLES_PREFIX,
VAR_KEY, SEP)
with self.assertLogs(secrets_manager_backend.log,
level="ERROR") as log_output:
self.assertIsNone(secrets_manager_backend.get_variable(VAR_KEY))
self.assertRegex(
log_output.output[0],
f"GCP API Call Error \\(NotFound\\): Secret ID {secret_id} not found"
)
def test_get_conn_uri_non_existent_key(self, mock_client_callable,
mock_get_creds):
mock_get_creds.return_value = CREDENTIALS, PROJECT_ID
mock_client = mock.MagicMock()
mock_client_callable.return_value = mock_client
# The requested secret id or secret version does not exist
mock_client.access_secret_version.side_effect = NotFound('test-msg')
connections_prefix = "airflow/connections"
secrets_manager_backend = CloudSecretsManagerBackend(
connections_prefix=connections_prefix)
with self.assertLogs(secrets_manager_backend.log,
level="ERROR") as log_output:
self.assertIsNone(
secrets_manager_backend.get_conn_uri(conn_id=CONN_ID))
self.assertEqual(
[], secrets_manager_backend.get_connections(conn_id=CONN_ID))
self.assertRegex(
log_output.output[0],
f"GCP API Call Error \\(NotFound\\): Secret ID {connections_prefix}/{CONN_ID} not found"
)
def test_get_connections(self, mock_get_uri, mock_get_creds):
mock_get_creds.return_value = CREDENTIALS, PROJECT_ID
mock_get_uri.return_value = CONN_URI
conns = CloudSecretsManagerBackend().get_connections(conn_id=CONN_ID)
self.assertIsInstance(conns, list)
self.assertIsInstance(conns[0], Connection)
def test_is_valid_prefix_and_sep(self, _, prefix, sep, is_valid):
backend = CloudSecretsManagerBackend()
backend.connections_prefix = prefix
backend.sep = sep
self.assertEqual(backend._is_valid_prefix_and_sep(), is_valid)
def test_raise_exception_with_invalid_prefix_sep(self, _, prefix, sep):
with self.assertRaises(AirflowException):
CloudSecretsManagerBackend(connections_prefix=prefix, sep=sep)
def test_default_valid_and_sep(self):
backend = CloudSecretsManagerBackend()
self.assertTrue(backend._is_valid_prefix_and_sep())
