def test_should_run_continue(load_dag, airflow_session, storage):
with mock.patch(
'google.cloud.storage.client.Client._connection',
new_callable=mock.PropertyMock,
) as client_mock:
client_mock.return_value = storage.mock_connection(
{'items': [{
'name': 'bar.csv'
}]})
start_date = timezone.datetime(2020, 1, 1)
dag_bag = load_dag('wrench_to_lake')
dag = dag_bag.get_dag('wrench_to_lake')
dag.clear()
should_run_task = dag.get_task('should_run')
assert isinstance(should_run_task, BranchPythonOperator)
should_run_ti = TaskInstance(should_run_task, start_date)
should_run_ti.run()
with airflow_session() as session:
tis = session.query(TaskInstance).filter(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.execution_date == start_date)
for ti in tis:
if ti.task_id == 'should_run':
assert ti.state == State.SUCCESS
elif ti.task_id == 'schedule_df_wrench_to_lake':
assert ti.state != State.SKIPPED
def test_get_dated_main_runner_handles_zero_shift():
dag = DAG(dag_id='test_dag',
start_date=datetime.strptime('2019-01-01', '%Y-%m-%d'))
execution_date = datetime.strptime('2019-01-01',
'%Y-%m-%d').replace(tzinfo=timezone.utc)
main_func = PickleMock()
runner = op_util.get_dated_main_runner_operator(dag, main_func,
timedelta(minutes=1))
ti = TaskInstance(runner, execution_date)
ti.run(ignore_task_deps=True, ignore_ti_state=True, test_mode=True)
main_func.assert_called_with('2019-01-01')
def _run_task(self, ti: TaskInstance) -> bool:
self.log.debug("Executing task: %s", ti)
key = ti.key
try:
params = self.tasks_params.pop(ti.key, {})
ti.run(job_id=ti.job_id, **params)
self.change_state(key, State.SUCCESS)
return True
except Exception as e:
ti.set_state(State.FAILED)
self.change_state(key, State.FAILED)
self.log.exception("Failed to execute task: %s.", str(e))
return False
class TestHelloOperator(unittest.TestCase):
def setUp(self):
self.dag = DAG(TEST_DAG_ID,
schedule_interval='@daily',
default_args={'start_date': DEFAULT_DATE})
self.op = HelloOperator(dag=self.dag,
task_id='test_hello',
name='test_name')
self.ti = TaskInstance(task=self.op, execution_date=DEFAULT_DATE)
def test_execute_no_trigger(self):
self.ti.run(ignore_ti_state=True)
self.assertEqual(self.ti.state, State.SUCCESS)
def run_task(self, op):
ti = TaskInstance(task=op, execution_date=DEFAULT_DATE)
ti.run(ignore_ti_state=True)
