def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger_tasks, mock_sync): executor = BaseExecutor() executor.heartbeat() calls = [mock.call('executor.open_slots', mock.ANY), mock.call('executor.queued_tasks', mock.ANY), mock.call('executor.running_tasks', mock.ANY)] mock_stats_gauge.assert_has_calls(calls)
def test_use_nf_executor(self): executor = BaseExecutor() executor.set_use_nf(True) executor.change_state('key', State.RUNNING) executor.change_state('key', State.SUCCESS) events = self.client.list_all_events(1) self.assertEqual(2, len(events))
def test_try_adopt_task_instances(self): date = datetime.utcnow() start_date = datetime.utcnow() - timedelta(days=2) with DAG("test_try_adopt_task_instances"): task_1 = BaseOperator(task_id="task_1", start_date=start_date) task_2 = BaseOperator(task_id="task_2", start_date=start_date) task_3 = BaseOperator(task_id="task_3", start_date=start_date) key1 = TaskInstance(task=task_1, execution_date=date) key2 = TaskInstance(task=task_2, execution_date=date) key3 = TaskInstance(task=task_3, execution_date=date) tis = [key1, key2, key3] self.assertEqual(BaseExecutor().try_adopt_task_instances(tis), tis)
def test_get_event_buffer(self): executor = BaseExecutor() date = datetime.utcnow() try_number = 1 key1 = ("my_dag1", "my_task1", date, try_number) key2 = ("my_dag2", "my_task1", date, try_number) key3 = ("my_dag2", "my_task2", date, try_number) state = State.SUCCESS executor.event_buffer[key1] = state, None executor.event_buffer[key2] = state, None executor.event_buffer[key3] = state, None self.assertEqual(len(executor.get_event_buffer(("my_dag1", ))), 1) self.assertEqual(len(executor.get_event_buffer()), 2) self.assertEqual(len(executor.event_buffer), 0)