def test_run_query(self, mock_execute):
mock_execute.return_value = None
db.merge_conn(
Connection(conn_id=ADX_TEST_CONN_ID,
conn_type='azure_data_explorer',
host='https://help.kusto.windows.net',
extra=json.dumps({'auth_method': 'AAD_DEVICE'})))
hook = AzureDataExplorerHook(
azure_data_explorer_conn_id=ADX_TEST_CONN_ID)
hook.run_query('Database',
'Logs | schema',
options={'option1': 'option_value'})
properties = ClientRequestProperties()
properties.set_option('option1', 'option_value')
assert mock_execute.called_with('Database',
'Logs | schema',
properties=properties)
def test_conn_missing_cluster(self):
db.merge_conn(
Connection(conn_id=ADX_TEST_CONN_ID,
conn_type='azure_data_explorer',
login='******',
password='******',
extra=json.dumps({})))
with self.assertRaises(AirflowException) as e:
AzureDataExplorerHook(azure_data_explorer_conn_id=ADX_TEST_CONN_ID)
self.assertIn('Host connection option is required', str(e.exception))
def test_conn_method_aad_device(self, mock_init):
mock_init.return_value = None
db.merge_conn(
Connection(conn_id=ADX_TEST_CONN_ID,
conn_type='azure_data_explorer',
host='https://help.kusto.windows.net',
extra=json.dumps({'auth_method': 'AAD_DEVICE'})))
AzureDataExplorerHook(azure_data_explorer_conn_id=ADX_TEST_CONN_ID)
assert mock_init.called_with(
KustoConnectionStringBuilder.with_aad_device_authentication(
'https://help.kusto.windows.net'))
def test_conn_missing_cluster(self):
db.merge_conn(
Connection(
conn_id=ADX_TEST_CONN_ID,
conn_type='azure_data_explorer',
login='******',
password='******',
extra=json.dumps({}),
))
with pytest.raises(AirflowException) as ctx:
AzureDataExplorerHook(azure_data_explorer_conn_id=ADX_TEST_CONN_ID)
assert 'Host connection option is required' in str(ctx.value)
def test_conn_unknown_method(self):
db.merge_conn(
Connection(conn_id=ADX_TEST_CONN_ID,
conn_type='azure_data_explorer',
login='******',
password='******',
host='https://help.kusto.windows.net',
extra=json.dumps({'auth_method': 'AAD_OTHER'})))
with self.assertRaises(AirflowException) as e:
AzureDataExplorerHook(azure_data_explorer_conn_id=ADX_TEST_CONN_ID)
self.assertIn('Unknown authentication method: AAD_OTHER',
str(e.exception))
def test_conn_missing_method(self):
db.merge_conn(
Connection(conn_id=ADX_TEST_CONN_ID,
conn_type='azure_data_explorer',
login='******',
password='******',
host='https://help.kusto.windows.net',
extra=json.dumps({})))
with self.assertRaises(AirflowException) as e:
AzureDataExplorerHook(azure_data_explorer_conn_id=ADX_TEST_CONN_ID)
self.assertIn('missing required parameter: `auth_method`',
str(e.exception))
def test_conn_method_aad_creds(self, mock_init):
mock_init.return_value = None
db.merge_conn(
Connection(conn_id=ADX_TEST_CONN_ID,
conn_type='azure_data_explorer',
login='******',
password='******',
host='https://help.kusto.windows.net',
extra=json.dumps({
'tenant': 'tenant',
'auth_method': 'AAD_CREDS'
})))
AzureDataExplorerHook(azure_data_explorer_conn_id=ADX_TEST_CONN_ID)
assert mock_init.called_with(
KustoConnectionStringBuilder.with_aad_user_password_authentication(
'https://help.kusto.windows.net', 'client_id', 'client secret',
'tenant'))
def test_conn_method_aad_app_cert(self, mock_init):
mock_init.return_value = None
db.merge_conn(
Connection(conn_id=ADX_TEST_CONN_ID,
conn_type='azure_data_explorer',
login='******',
host='https://help.kusto.windows.net',
extra=json.dumps({
'tenant': 'tenant',
'auth_method': 'AAD_APP_CERT',
'certificate': 'PEM',
'thumbprint': 'thumbprint'
})))
AzureDataExplorerHook(azure_data_explorer_conn_id=ADX_TEST_CONN_ID)
assert mock_init.called_with(
KustoConnectionStringBuilder.
with_aad_application_certificate_authentication(
'https://help.kusto.windows.net', 'client_id', 'PEM',
'thumbprint', 'tenant'))
def get_hook(self) -> AzureDataExplorerHook:
"""Returns new instance of AzureDataExplorerHook"""
return AzureDataExplorerHook(self.azure_data_explorer_conn_id)
