def test_get_variable(self, mock_client):
mock_client.get_secret.return_value = mock.Mock(value='world')
backend = AzureKeyVaultBackend()
returned_uri = backend.get_variable('hello')
mock_client.get_secret.assert_called_with(
name='airflow-variables-hello')
assert 'world' == returned_uri
def test_get_variable_non_existent_key(self, mock_client):
"""
Test that if Variable key is not present,
AzureKeyVaultBackend.get_variables should return None
"""
mock_client.get_secret.side_effect = ResourceNotFoundError
backend = AzureKeyVaultBackend()
assert backend.get_variable('test_mysql') is None
def test_get_secret_value_not_found(self, mock_client):
"""
Test that if a non-existent secret returns None
"""
mock_client.get_secret.side_effect = ResourceNotFoundError
backend = AzureKeyVaultBackend()
assert (backend._get_secret(path_prefix=backend.connections_prefix,
secret_id='test_non_existent') is None)
def test_get_secret_value(self, mock_client):
"""
Test that get_secret returns the secret value
"""
mock_client.get_secret.return_value = mock.Mock(value='super-secret')
backend = AzureKeyVaultBackend()
secret_val = backend._get_secret('af-secrets', 'test_mysql_password')
mock_client.get_secret.assert_called_with(name='af-secrets-test-mysql-password')
self.assertEqual(secret_val, 'super-secret')
def test_config_prefix_none_value(self, mock_get_secret):
"""
Test that if Config prefix is None,
AzureKeyVaultBackend.get_config should return None
AzureKeyVaultBackend._get_secret should not be called
"""
kwargs = {'config_prefix': None}
backend = AzureKeyVaultBackend(**kwargs)
assert backend.get_config('test_mysql') is None
mock_get_secret.assert_not_called()
def test_variable_prefix_none_value(self, mock_get_secret):
"""
Test that if Variables prefix is None,
AzureKeyVaultBackend.get_variables should return None
AzureKeyVaultBackend._get_secret should not be called
"""
kwargs = {'variables_prefix': None}
backend = AzureKeyVaultBackend(**kwargs)
assert backend.get_variable('hello') is None
mock_get_secret.assert_not_called()
def test_get_conn_uri_non_existent_key(self, mock_client):
"""
Test that if the key with connection ID is not present,
AzureKeyVaultBackend.get_connections should return None
"""
conn_id = 'test_mysql'
mock_client.get_secret.side_effect = ResourceNotFoundError
backend = AzureKeyVaultBackend(vault_url="https://example-akv-resource-name.vault.azure.net/")
self.assertIsNone(backend.get_conn_uri(conn_id=conn_id))
self.assertEqual([], backend.get_connections(conn_id=conn_id))
def test_connection_prefix_none_value(self, mock_get_secret):
"""
Test that if Connections prefix is None,
AzureKeyVaultBackend.get_connections should return None
AzureKeyVaultBackend._get_secret should not be called
"""
kwargs = {'connections_prefix': None}
backend = AzureKeyVaultBackend(**kwargs)
self.assertIsNone(backend.get_conn_uri('test_mysql'))
mock_get_secret._get_secret.assert_not_called()
def test_get_conn_uri(self, mock_secret_client, mock_azure_cred):
mock_cred = mock.Mock()
mock_sec_client = mock.Mock()
mock_azure_cred.return_value = mock_cred
mock_secret_client.return_value = mock_sec_client
mock_sec_client.get_secret.return_value = mock.Mock(
value='postgresql://*****:*****@host:5432/airflow')
backend = AzureKeyVaultBackend(
vault_url="https://example-akv-resource-name.vault.azure.net/")
returned_uri = backend.get_conn_uri(conn_id='hi')
mock_secret_client.assert_called_once_with(
credential=mock_cred,
vault_url='https://example-akv-resource-name.vault.azure.net/')
assert returned_uri == 'postgresql://*****:*****@host:5432/airflow'
def test_get_connections(self, mock_get_uri):
mock_get_uri.return_value = 'scheme://*****:*****@host:100'
conn_list = AzureKeyVaultBackend().get_connections('fake_conn')
conn = conn_list[0]
assert conn.host == 'host'