class Test(unittest.TestCase):
def setUp(self):
session = aiohttp.ClientSession(
raise_for_status=True, timeout=aiohttp.ClientTimeout(total=60))
self.client = BatchClient(session, url=os.environ.get('BATCH_URL'))
def tearDown(self):
loop = asyncio.get_event_loop()
loop.run_until_complete(self.client.close())
def test_job(self):
async def f():
b = self.client.create_batch()
j = b.create_job('alpine', ['echo', 'test'])
await b.submit()
status = await j.wait()
self.assertTrue('attributes' not in status)
self.assertEqual(status['state'], 'Success')
self.assertEqual(status['exit_code']['main'], 0)
self.assertEqual(await j.log(), {'main': 'test\n'})
self.assertTrue(await j.is_complete())
loop = asyncio.get_event_loop()
loop.run_until_complete(f())
async def test_bad_token():
token = session_id_encode_to_str(secrets.token_bytes(32))
bc = BatchClient('test', _token=token)
try:
b = bc.create_batch()
j = b.create_job(DOCKER_ROOT_IMAGE, ['false'])
await b.submit()
assert False, j
except aiohttp.ClientResponseError as e:
assert e.status == 401, e
finally:
await bc.close()
async def random_billing_project_name(
dev_client: BatchClient) -> AsyncGenerator[str, Any]:
billing_project_prefix = get_billing_project_prefix()
name = f'{billing_project_prefix}_{secret_alnum_string(5)}'
try:
yield name
finally:
try:
r = await dev_client.get_billing_project(name)
except aiohttp.ClientResponseError as e:
assert e.status == 403, e
else:
assert r['status'] != 'deleted', r
try:
async for batch in dev_client.list_batches(
f'billing_project:{name}'):
await batch.delete()
finally:
try:
if r['status'] == 'open':
await dev_client.close_billing_project(name)
finally:
if r['status'] != 'deleted':
await dev_client.delete_billing_project(name)
async def on_startup(app):
app['gh_client_session'] = aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=5))
app['github_client'] = gh_aiohttp.GitHubAPI(app['gh_client_session'],
'ci',
oauth_token=oauth_token)
app['batch_client'] = BatchClient('ci')
app['dbpool'] = await create_database_pool()
app['task_manager'] = aiotools.BackgroundTaskManager()
app['task_manager'].ensure_future(update_loop(app))
async def delete_all_test_billing_projects():
billing_project_prefix = get_billing_project_prefix()
bc = BatchClient(None, token_file=os.environ['HAIL_TEST_DEV_TOKEN_FILE'])
try:
for project in await bc.list_billing_projects():
if project['billing_project'].startswith(billing_project_prefix):
try:
await bc.close_billing_project(project['billing_project'])
finally:
await bc.delete_billing_project(project['billing_project'])
finally:
await bc.close()
async def test_job(client: BatchClient):
b = client.create_batch()
j = b.create_job(DOCKER_ROOT_IMAGE, ['echo', 'test'])
b = await b.submit()
status = await j.wait()
assert 'attributes' not in status, str((status, await b.debug_info()))
assert status['state'] == 'Success', str((status, await b.debug_info()))
assert j._get_exit_code(status, 'main') == 0, str((status, await
b.debug_info()))
job_log = await j.log()
assert job_log['main'] == 'test\n', str((job_log, await b.debug_info()))
assert await j.is_complete(), str(await b.debug_info())
async def on_startup(app):
app['client_session'] = aiohttp.ClientSession(
raise_for_status=True,
timeout=aiohttp.ClientTimeout(total=60))
app['github_client'] = gh_aiohttp.GitHubAPI(app['client_session'], 'ci', oauth_token=oauth_token)
app['batch_client'] = BatchClient(app['client_session'], url=os.environ.get('BATCH_SERVER_URL'))
with open('/ci-user-secret/sql-config.json', 'r') as f:
config = json.loads(f.read().strip())
app['dbpool'] = await aiomysql.create_pool(host=config['host'],
port=config['port'],
db=config['db'],
user=config['user'],
password=config['password'],
charset='utf8',
cursorclass=aiomysql.cursors.DictCursor,
autocommit=True)
asyncio.ensure_future(update_loop(app))
async def on_startup(app):
app['gh_client_session'] = aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=5))
app['github_client'] = gh_aiohttp.GitHubAPI(app['gh_client_session'], 'ci')
app['batch_client'] = BatchClient('ci')
def setUp(self):
session = aiohttp.ClientSession(
raise_for_status=True, timeout=aiohttp.ClientTimeout(total=60))
self.client = BatchClient(session, url=os.environ.get('BATCH_URL'))
def setUp(self):
self.client = async_to_blocking(BatchClient())
async def client():
bc = BatchClient('test')
yield bc
await bc.close()
async def dev_client() -> AsyncGenerator[BatchClient, Any]:
bc = BatchClient('billing-project-not-needed-but-required-by-BatchClient',
token_file=os.environ['HAIL_TEST_DEV_TOKEN_FILE'])
yield bc
await bc.close()
async def factory(project):
bc = BatchClient(project, token_file=os.environ['HAIL_TEST_TOKEN_FILE'])
_bcs.append(bc)
return bc
async def dev_client():
bc = BatchClient(None, token_file=os.environ['HAIL_TEST_DEV_TOKEN_FILE'])
yield bc
await bc.close()