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'