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)
assert backend.get_conn_uri('test_mysql') 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_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'