State.SCHEDULED,
}
# Context to get the dependencies that need to be met in order for a task instance to be
# set to 'scheduled' state.
SCHEDULED_DEPS = {
RunnableExecDateDep(),
ValidStateDep(SCHEDULEABLE_STATES),
TaskNotRunningDep(),
}
# Dependencies that if met, task instance should be re-queued.
REQUEUEABLE_DEPS = {
DagTISlotsAvailableDep(),
TaskConcurrencyDep(),
PoolSlotsAvailableDep(),
}
# Dependencies that need to be met for a given task instance to be set to 'RUNNING' state.
RUNNING_DEPS = {
RunnableExecDateDep(),
ValidStateDep(RUNNABLE_STATES),
DagTISlotsAvailableDep(),
TaskConcurrencyDep(),
PoolSlotsAvailableDep(),
TaskNotRunningDep(),
}
BACKFILL_QUEUED_DEPS = {
RunnableExecDateDep(),
ValidStateDep(BACKFILL_QUEUEABLE_STATES),
def test_task_with_nonexistent_pool(self):
ti = Mock(pool='nonexistent_pool', pool_slots=1)
self.assertFalse(PoolSlotsAvailableDep().is_met(ti=ti))
def test_running_pooled_task_pass(self, mock_open_slots):
for state in EXECUTION_STATES:
ti = Mock(pool='test_pool', state=state, pool_slots=1)
self.assertTrue(PoolSlotsAvailableDep().is_met(ti=ti))
def test_pooled_task_pass(self, mock_open_slots):
ti = Mock(pool='test_pool', pool_slots=1)
self.assertTrue(PoolSlotsAvailableDep().is_met(ti=ti))
def test_pooled_task_reached_concurrency(self, mock_open_slots):
ti = Mock(pool='test_pool', pool_slots=1)
self.assertFalse(PoolSlotsAvailableDep().is_met(ti=ti))
def test_running_pooled_task_pass(self, mock_open_slots):
for state in STATES_TO_COUNT_AS_RUNNING:
ti = Mock(pool='test_pool', state=state, pool_slots=1)
self.assertTrue(PoolSlotsAvailableDep().is_met(ti=ti))