class TestAzureDataExplorerQueryOperator(unittest.TestCase): def setUp(self): args = {'owner': 'airflow', 'start_date': DEFAULT_DATE, 'provide_context': True} self.dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args, schedule_interval='@once') self.operator = AzureDataExplorerQueryOperator(dag=self.dag, **MOCK_DATA) def test_init(self): self.assertEqual(self.operator.task_id, MOCK_DATA['task_id']) self.assertEqual(self.operator.query, MOCK_DATA['query']) self.assertEqual(self.operator.database, MOCK_DATA['database']) self.assertEqual(self.operator.azure_data_explorer_conn_id, 'azure_data_explorer_default') @mock.patch.object(AzureDataExplorerHook, 'run_query', return_value=MockResponse()) @mock.patch.object(AzureDataExplorerHook, 'get_conn') def test_run_query(self, mock_conn, mock_run_query): self.operator.execute(None) mock_run_query.assert_called_once_with( MOCK_DATA['query'], MOCK_DATA['database'], MOCK_DATA['options'] ) @mock.patch.object(AzureDataExplorerHook, 'run_query', return_value=MockResponse()) @mock.patch.object(AzureDataExplorerHook, 'get_conn') def test_xcom_push_and_pull(self, mock_conn, mock_run_query): ti = TaskInstance(task=self.operator, execution_date=timezone.utcnow()) ti.run() self.assertEqual(ti.xcom_pull(task_ids=MOCK_DATA['task_id']), str(MOCK_RESULT))
def setUp(self): args = { 'owner': 'airflow', 'start_date': DEFAULT_DATE, 'provide_context': True } self.dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args, schedule_interval='@once') self.operator = AzureDataExplorerQueryOperator(dag=self.dag, **MOCK_DATA)
from airflow import DAG from airflow.providers.microsoft.azure.operators.adx import AzureDataExplorerQueryOperator from datetime import datetime, timedelta adx_query = '''StormEvents | sort by StartTime desc | take 10''' # Default settings applied to all tasks default_args = { 'owner': 'airflow', 'depends_on_past': False, 'email_on_failure': False, 'email_on_retry': False, 'retries': 0, 'retry_delay': timedelta(minutes=1) } with DAG('azure_data_explorer', start_date=datetime(2020, 12, 1), max_active_runs=1, schedule_interval='@daily', default_args=default_args, catchup=False) as dag: opr_adx_query = AzureDataExplorerQueryOperator( task_id='adx_query', query=adx_query, database='storm_demo', azure_data_explorer_conn_id='adx')