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)