def test_python_callable_keyword_arguments_are_templatized(self): """Test PythonSensor op_kwargs are templatized""" recorded_calls = [] task = PythonSensor( task_id='python_sensor', timeout=0.01, poke_interval=0.3, # 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_type=DagRunType.MANUAL, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING) with self.assertRaises(AirflowSensorTimeout): task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # 2 calls: first: at start, second: before timeout self.assertEqual(2, 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_python_callable_arguments_are_templatized(self): """Test PythonSensor 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 = PythonSensor( task_id='python_sensor', timeout=0.01, poke_interval=0.3, # 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_type=DagRunType.MANUAL, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, state=State.RUNNING, ) with pytest.raises(AirflowSensorTimeout): task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) ds_templated = DEFAULT_DATE.date().isoformat() # 2 calls: first: at start, second: before timeout assert 2 == len(recorded_calls) self._assert_calls_equal( recorded_calls[0], Call( 4, date(2019, 1, 1), f"dag {self.dag.dag_id} ran on {ds_templated}.", Named(ds_templated, 'unchanged'), ), )