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_default_pool_open_slots(self):
set_default_pool_slots(5)
self.assertEqual(5, Pool.get_default_pool().open_slots())
dag = DAG(
dag_id='test_default_pool_open_slots',
start_date=DEFAULT_DATE, )
t1 = DummyOperator(task_id='dummy1', dag=dag)
t2 = DummyOperator(task_id='dummy2', dag=dag)
ti1 = TI(task=t1, execution_date=DEFAULT_DATE)
ti2 = TI(task=t2, 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()
self.assertEqual(3, Pool.get_default_pool().open_slots())
def setUp(self):
clear_db_pools()
self.pools = [Pool.get_default_pool()]
for i in range(self.USER_POOL_COUNT):
name = f'experimental_{i + 1}'
pool = models.Pool(
pool=name,
slots=i,
description=name,
)
self.pools.append(pool)
with create_session() as session:
session.add_all(self.pools)
def setUp(self):
self.session = settings.Session()
clear_db_pools()
self.pools = [Pool.get_default_pool()]
for i in range(self.USER_POOL_COUNT):
name = 'experimental_%s' % (i + 1)
pool = models.Pool(
pool=name,
slots=i,
description=name,
)
self.session.add(pool)
self.pools.append(pool)
self.session.commit()
