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')