def test_python_callable_arguments_are_templatized(self): """Test PythonOperator op_args are templatized""" recorded_calls = [] # Create a named tuple and ensure it is still preserved # after the rendering is done Named = namedtuple('Named', ['var1', 'var2']) named_tuple = Named('{{ ds }}', 'unchanged') task = PythonOperator( task_id='python_operator', # a Mock instance cannot be used as a callable function or test fails with a # TypeError: Object of type Mock is not JSON serializable python_callable=build_recording_function(recorded_calls), op_args=[ 4, date(2019, 1, 1), "dag {{dag.dag_id}} ran on {{ds}}.", named_tuple ], dag=self.dag) self.dag.create_dagrun(run_id='manual__' + DEFAULT_DATE.isoformat(), execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING) task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) ds_templated = DEFAULT_DATE.date().isoformat() self.assertEqual(1, len(recorded_calls)) self._assert_calls_equal( recorded_calls[0], Call(4, date(2019, 1, 1), "dag {} ran on {}.".format(self.dag.dag_id, ds_templated), Named(ds_templated, 'unchanged')))
def test_python_callable_keyword_arguments_are_templatized(self): """Test PythonOperator op_kwargs are templatized""" recorded_calls = [] task = PythonOperator( task_id='python_operator', # a Mock instance cannot be used as a callable function or test fails with a # TypeError: Object of type Mock is not JSON serializable python_callable=build_recording_function(recorded_calls), op_kwargs={ 'an_int': 4, 'a_date': date(2019, 1, 1), 'a_templated_string': "dag {{dag.dag_id}} ran on {{ds}}." }, dag=self.dag) self.dag.create_dagrun(run_id='manual__' + DEFAULT_DATE.isoformat(), execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING) task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) self.assertEqual(1, len(recorded_calls)) self._assert_calls_equal( recorded_calls[0], Call(an_int=4, a_date=date(2019, 1, 1), a_templated_string="dag {} ran on {}.".format( self.dag.dag_id, DEFAULT_DATE.date().isoformat())))
def test_conflicting_kwargs(self): self.dag.create_dagrun( run_type=DagRunType.MANUAL, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING, external_trigger=False, ) # dag is not allowed since it is a reserved keyword def func(dag): # An ValueError should be triggered since we're using dag as a # reserved keyword raise RuntimeError("Should not be triggered, dag: {}".format(dag)) python_operator = PythonOperator( task_id='python_operator', op_args=[1], python_callable=func, dag=self.dag ) with self.assertRaises(ValueError) as context: python_operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) self.assertTrue('dag' in context.exception, "'dag' not found in the exception")
def test_python_operator_run(self): """Tests that the python callable is invoked on task run.""" task = PythonOperator(python_callable=self.do_run, task_id='python_operator', dag=self.dag) self.assertFalse(self.is_run()) task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) self.assertTrue(self.is_run())
def test_timeout(self): op = PythonOperator( task_id='test_timeout', execution_timeout=timedelta(seconds=1), python_callable=lambda: sleep(5), dag=self.dag, ) with pytest.raises(AirflowTaskTimeout): op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
def test_python_op(self): def test_py_op(templates_dict, ds, **kwargs): if not templates_dict['ds'] == ds: raise Exception("failure") op = PythonOperator( task_id='test_py_op', python_callable=test_py_op, templates_dict={'ds': "{{ ds }}"}, dag=self.dag) op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
def test_python_op(self): def test_py_op(templates_dict, ds, **kwargs): if not templates_dict['ds'] == ds: raise Exception("failure") op = PythonOperator(task_id='test_py_op', python_callable=test_py_op, templates_dict={'ds': "{{ ds }}"}, dag=self.dag) self.dag.create_dagrun(run_type=DagRunType.MANUAL, state=State.RUNNING, execution_date=DEFAULT_DATE) op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
def test_echo_env_variables(self): """ Test that env variables are exported correctly to the python callback in the task. """ self.dag.create_dagrun( run_id='manual__' + DEFAULT_DATE.isoformat(), execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING, external_trigger=False, ) op = PythonOperator(task_id='hive_in_python_op', dag=self.dag, python_callable=self._env_var_check_callback) op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
def test_context_with_conflicting_op_args(self): self.dag.create_dagrun( run_id='manual__' + DEFAULT_DATE.isoformat(), execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING, external_trigger=False, ) def func(custom, dag): self.assertEqual(1, custom, "custom should be 1") self.assertIsNotNone(dag, "dag should be set") python_operator = PythonOperator(task_id='python_operator', op_kwargs={'custom': 1}, python_callable=func, dag=self.dag) python_operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
def test_context_with_kwargs(self): self.dag.create_dagrun( run_type=DagRunType.MANUAL, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING, external_trigger=False, ) def func(**context): # check if context is being set assert len(context) > 0, "Context has not been injected" python_operator = PythonOperator(task_id='python_operator', op_kwargs={'custom': 1}, python_callable=func, dag=self.dag) python_operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
def test_context_with_conflicting_op_args(self): self.dag.create_dagrun( run_type=DagRunType.MANUAL, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING, external_trigger=False, ) def func(custom, dag): assert 1 == custom, "custom should be 1" assert dag is not None, "dag should be set" python_operator = PythonOperator(task_id='python_operator', op_kwargs={'custom': 1}, python_callable=func, dag=self.dag) python_operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
def test_context_with_kwargs(self): self.dag.create_dagrun( run_id='manual__' + DEFAULT_DATE.isoformat(), execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING, external_trigger=False, ) def func(**context): # check if context is being set self.assertGreater(len(context), 0, "Context has not been injected") python_operator = PythonOperator(task_id='python_operator', op_kwargs={'custom': 1}, python_callable=func, dag=self.dag) python_operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
def test_provide_context_does_not_fail(self): """ ensures that provide_context doesn't break dags in 2.0 """ self.dag.create_dagrun( run_type=DagRunType.MANUAL, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING, external_trigger=False, ) def func(custom, dag): assert 1 == custom, "custom should be 1" assert dag is not None, "dag should be set" python_operator = PythonOperator( task_id='python_operator', op_kwargs={'custom': 1}, python_callable=func, provide_context=True, dag=self.dag, ) python_operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
def test_get_context_in_old_style_context_task(self): with DAG(dag_id="edge_case_context_dag", default_args=DEFAULT_ARGS): op = PythonOperator(python_callable=get_all_the_context, task_id="get_all_the_context") op.run(ignore_first_depends_on_past=True, ignore_ti_state=True)