def test_default_pool_open_slots(self):
set_default_pool_slots(5)
assert 5 == Pool.get_default_pool().open_slots()
dag = DAG(
dag_id='test_default_pool_open_slots',
start_date=DEFAULT_DATE,
)
op1 = DummyOperator(task_id='dummy1', dag=dag)
op2 = DummyOperator(task_id='dummy2', dag=dag, pool_slots=2)
ti1 = TI(task=op1, execution_date=DEFAULT_DATE)
ti2 = TI(task=op2, execution_date=DEFAULT_DATE)
ti1.state = State.RUNNING
ti2.state = State.QUEUED
session = settings.Session
session.add(ti1)
session.add(ti2)
session.commit()
session.close()
assert 2 == Pool.get_default_pool().open_slots()
assert {
"default_pool": {
"open": 2,
"queued": 2,
"total": 5,
"running": 1,
}
} == Pool.slots_stats()
def test_open_slots(self):
pool = Pool(pool='test_pool', slots=5)
dag = DAG(
dag_id='test_open_slots',
start_date=DEFAULT_DATE,
)
op1 = DummyOperator(task_id='dummy1', dag=dag, pool='test_pool')
op2 = DummyOperator(task_id='dummy2', dag=dag, pool='test_pool')
ti1 = TI(task=op1, execution_date=DEFAULT_DATE)
ti2 = TI(task=op2, execution_date=DEFAULT_DATE)
ti1.state = State.RUNNING
ti2.state = State.QUEUED
session = settings.Session
session.add(pool)
session.add(ti1)
session.add(ti2)
session.commit()
session.close()
self.assertEqual(3, pool.open_slots()) # pylint: disable=no-value-for-parameter
self.assertEqual(1, pool.running_slots()) # pylint: disable=no-value-for-parameter
self.assertEqual(1, pool.queued_slots()) # pylint: disable=no-value-for-parameter
self.assertEqual(2, pool.occupied_slots()) # pylint: disable=no-value-for-parameter
self.assertEqual(
{
"default_pool": {
"open": 128,
"queued": 0,
"total": 128,
"running": 0,
},
"test_pool": {
"open": 3,
"queued": 1,
"running": 1,
"total": 5,
},
},
pool.slots_stats(),
)
