def setUp(self):
clear_db_runs()
clear_db_pools()
clear_db_dags()
clear_db_sla_miss()
clear_db_errors()
clear_db_event_model()
def tearDown(self):
db.clear_db_pools()
with create_session() as session:
session.query(TaskFail).delete()
session.query(TaskReschedule).delete()
session.query(models.TaskInstance).delete()
session.query(models.DagRun).delete()
def test_delete_default_pool(self):
clear_db_pools()
response = self.client.delete('/api/experimental/pools/default_pool', )
self.assertEqual(response.status_code, 400)
self.assertEqual(
json.loads(response.data.decode('utf-8'))['error'],
"default_pool cannot be deleted")
def test_success(self):
token = "Basic " + b64encode(b"test:test").decode()
clear_db_pools()
with self.app.test_client() as test_client:
response = test_client.get("/api/experimental/pools",
headers={"Authorization": token})
assert current_user.email == "*****@*****.**"
assert response.status_code == 200
def test_success_using_email(self):
clear_db_pools()
with self.app.test_client() as test_client:
response = test_client.get(
"/api/experimental/pools", environ_overrides={'REMOTE_USER': "******"}
)
self.assertEqual("*****@*****.**", current_user.email)
self.assertEqual(200, response.status_code)
self.assertIn("Default pool", str(response.json))
def test_success_using_email(self):
clear_db_pools()
with self.app.test_client() as test_client:
response = test_client.get(
"/api/experimental/pools",
environ_overrides={'REMOTE_USER': "******"})
assert "*****@*****.**" == current_user.email
assert 200 == response.status_code
assert "Default pool" in str(response.json)
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_attrs(self, _setup_attrs_base):
clear_db_pools()
self.pools = [Pool.get_default_pool()]
for i in range(self.USER_POOL_COUNT):
name = f'experimental_{i + 1}'
pool = Pool(
pool=name,
slots=i,
description=name,
)
self.session.add(pool)
self.pools.append(pool)
self.session.commit()
self.pool = self.pools[-1]
def test_experimental_api(self):
with self.app.test_client() as test_client:
response = test_client.get("/api/experimental/pools", headers={"Authorization": "Basic"})
assert response.status_code == 401
assert response.headers["WWW-Authenticate"] == "Basic"
assert response.data == b'Unauthorized'
clear_db_pools()
response = test_client.get(
"/api/experimental/pools",
headers={"Authorization": "Basic " + b64encode(b"test:test").decode()},
)
assert response.status_code == 200
assert response.json[0]["pool"] == 'default_pool'
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()
def test_run_pooling_task(self):
"""
test that running a task in an existing pool update task state as SUCCESS.
"""
dag = models.DAG(dag_id='test_run_pooling_task')
task = DummyOperator(task_id='test_run_pooling_task_op', dag=dag,
pool='test_pool', owner='airflow',
start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))
ti = TI(
task=task, execution_date=timezone.utcnow())
ti.run()
db.clear_db_pools()
self.assertEqual(ti.state, State.SUCCESS)
def setUp(self):
super().setUp()
clear_db_pools()
self.pools = [Pool.get_default_pool()]
for i in range(self.USER_POOL_COUNT):
name = f'experimental_{i + 1}'
pool = Pool(
pool=name,
slots=i,
description=name,
)
self.session.add(pool)
self.pools.append(pool)
self.session.commit()
self.pool = self.pools[-1]
def test_success(self, mock_verify_token):
clear_db_pools()
mock_verify_token.return_value = {
"iss": "accounts.google.com",
"email_verified": True,
"email": "*****@*****.**",
}
with self.app.test_client() as test_client:
response = test_client.get(
"/api/experimental/pools",
headers={"Authorization": "bearer JWT_TOKEN"})
self.assertEqual("*****@*****.**", current_user.email)
self.assertEqual(200, response.status_code)
self.assertIn("Default pool", str(response.json))
def setUp(self):
super(TestPoolApiExperimental, self).setUp()
app = application.create_app(testing=True)
self.app = app.test_client()
self.session = Session()
clear_db_pools()
self.pools = [Pool.get_default_pool()]
for i in range(self.USER_POOL_COUNT):
name = 'experimental_%s' % (i + 1)
pool = Pool(
pool=name,
slots=i,
description=name,
)
self.session.add(pool)
self.pools.append(pool)
self.session.commit()
self.pool = self.pools[-1]
def test_success(self):
token = "Basic " + b64encode(b"test:test").decode()
clear_db_pools()
with self.app.test_client() as test_client:
response = test_client.get("/api/v1/pools", headers={"Authorization": token})
assert current_user.email == "*****@*****.**"
assert response.status_code == 200
assert response.json == {
"pools": [
{
"name": "default_pool",
"slots": 128,
"occupied_slots": 0,
"running_slots": 0,
"queued_slots": 0,
"open_slots": 128,
},
],
"total_entries": 1,
}
def tearDown(self):
clear_db_pools()
super(TestLocalClient, self).tearDown()
def setUp(self):
db.clear_db_pools()
with create_session() as session:
test_pool = Pool(pool='test_pool', slots=1)
session.add(test_pool)
session.commit()
def setUp(self):
super().setUp()
clear_db_pools()
self.client = Client(api_base_url=None, auth=None)
def tearDown(self):
clear_db_pools()
super().tearDown()
def clean_db():
clear_db_runs()
clear_db_pools()
def tearDown(self):
clear_db_pools()
super(TestLocalClient, self).tearDown()
def setUp(self) -> None:
clear_db_pools()
def tearDown(self):
clear_db_runs()
clear_db_pools()
def setUp(self):
clear_db_runs()
clear_db_pools()
def setUp(self):
clear_db_runs()
clear_db_pools()
self.parser = cli_parser.get_parser()
def tearDown(self) -> None:
clear_db_pools()
def setUp(self) -> None:
self.client = self.app.test_client() # type:ignore
super().setUp()
clear_db_pools()
def setUp(self):
super().setUp()
clear_db_pools()
self.client = Client(api_base_url=None, auth=None)
def tearDown(self):
clear_db_pools()
super().tearDown()
def tearDown(self):
db.clear_db_pools()
def tearDown(self):
clear_db_runs()
clear_db_pools()
def setup_attrs(self, configured_app) -> None:
self.app = configured_app
self.client = self.app.test_client() # type:ignore
clear_db_pools()
def setUp(self):
clear_db_runs()
clear_db_pools()
self.parser = cli.CLIFactory.get_parser()
def teardown_method(self) -> None:
clear_db_pools()
