async def test_api_function_metaclass_async():
async with AsyncSession() as session:
Dummy = type('DummyFunction', (BaseFunction, ), {
**DummyFunction.__dict__,
'session': session,
})
assert Dummy.session is session
assert Dummy().session is session
assert await Dummy.get_or_create() == 'created'
assert await Dummy().calculate() == 'done'
assert await Dummy.get_or_create() == 'created'
assert await Dummy().calculate() == 'done'
async with AsyncSession() as session:
Dummy = type('DummyFunction', (BaseFunction, ), {
**DummyFunction.__dict__,
'session': session,
})
assert Dummy.session is session
assert Dummy().session is session
assert await Dummy.get_or_create() == 'created'
assert await Dummy().calculate() == 'done'
assert await Dummy.get_or_create() == 'created'
assert await Dummy().calculate() == 'done'
async def test_fetch_timeout_async(dummy_endpoint):
with aioresponses() as m:
async with AsyncSession() as session:
m.post(dummy_endpoint, exception=asyncio.TimeoutError())
rqst = Request(session, 'POST', '/')
with pytest.raises(asyncio.TimeoutError):
async with rqst.fetch():
pass
async def test_fetch_client_error_async(dummy_endpoint):
with aioresponses() as m:
async with AsyncSession() as session:
m.post(dummy_endpoint, exception=aiohttp.ClientConnectionError())
rqst = Request(session, 'POST', '/')
with pytest.raises(BackendClientError):
async with rqst.fetch():
pass
async def test_restart_kernel_raises_err_with_abnormal_status():
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.afetch.return_value = asynctest.MagicMock(status=400)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj):
with pytest.raises(BackendAPIError):
await session.Kernel('mykernel').restart()
async def test_fetch_cancellation_async(self, dummy_endpoint):
with aioresponses() as m:
async with AsyncSession() as session:
m.post(dummy_endpoint,
exception=asyncio.CancelledError())
rqst = Request(session, 'POST', '/')
with pytest.raises(asyncio.CancelledError):
async with rqst.fetch():
pass
async def test_execute_code_raises_err_with_abnormal_status():
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.afetch.return_value = asynctest.MagicMock(status=400)
run_id = token_hex(8)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj):
with pytest.raises(BackendAPIError):
await session.Kernel('mykernel').execute(run_id, 'hello')
async def test_fetch_cancellation_async(dummy_endpoint):
# It seems that aiohttp swallows asyncio.CancelledError
with aioresponses() as m:
async with AsyncSession() as session:
m.post(dummy_endpoint,
exception=asyncio.CancelledError())
rqst = Request('POST', '/')
with pytest.raises(asyncio.CancelledError):
async with rqst.fetch():
pass
async def test_restart_kernel_url(mocker):
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.afetch.return_value = asynctest.MagicMock(status=204)
kernel_id = token_hex(12)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel(kernel_id).restart()
mock_req_cls.assert_called_once_with(
session, 'PATCH', '/kernel/{}'.format(kernel_id))
async def test_execute_code_url(mocker):
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.afetch.return_value = asynctest.MagicMock(status=200)
kernel_id = token_hex(12)
run_id = token_hex(8)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel(kernel_id).execute(run_id, 'hello')
mock_req_cls.assert_called_once_with(
session, 'POST', '/kernel/{}'.format(kernel_id),
{'mode': 'query', 'runId': run_id, 'code': 'hello'})
async def test_create_kernel_url(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=201,
json=AsyncMock())
mock_req_cls = mocker.patch('ai.backend.client.func.session.Request',
return_value=mock_req_obj)
async with AsyncSession() as session:
await session.ComputeSession.get_or_create('python:3.6-ubuntu18.04')
prefix = get_naming(session.api_version, 'path')
mock_req_cls.assert_called_once_with(session, 'POST', f'/{prefix}')
mock_req_obj.fetch.assert_called_once_with()
mock_req_obj.fetch.return_value.json.assert_awaited_once_with()
async def test_stream_pty_raises_error_with_abnormal_status(mocker):
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_exception = aiohttp.ClientResponseError(
None, None, code=400,
message='emulated-handshake-error')
mock_req_obj.connect_websocket = \
asynctest.MagicMock(side_effect=mock_exception)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj):
with pytest.raises(BackendClientError):
await session.Kernel('mykernel').stream_pty()
async def test_async_auth():
random_msg = uuid.uuid4().hex
async with AsyncSession() as sess:
request = Request('GET', '/auth')
request.set_json({
'echo': random_msg,
})
async with request.fetch() as resp:
assert resp.status == 200
data = await resp.json()
assert data['authorized'] == 'yes'
assert data['echo'] == random_msg
async def test_response_async(defconfig, dummy_endpoint):
body = b'{"test": 5678}'
with aioresponses() as m:
m.post(
dummy_endpoint + 'function', status=200, body=body,
headers={'Content-Type': 'application/json',
'Content-Length': str(len(body))},
)
async with AsyncSession(config=defconfig) as session:
rqst = Request('POST', '/function')
async with rqst.fetch() as resp:
assert await resp.text() == '{"test": 5678}'
assert await resp.json() == {'test': 5678}
async def test_create_kernel_return_id_only(self):
return_value = {'kernelId': 'mock_kernel_id'}
mock_json_coro = asynctest.CoroutineMock(return_value=return_value)
mock_req_obj = asynctest.MagicMock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=201,
json=mock_json_coro)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj):
k = await session.Kernel.get_or_create('python:3.6-ubuntu18.04'
)
assert k.kernel_id == return_value['kernelId']
async def test_destroy_kernel_url(self, mocker):
mock_req_obj = asynctest.MagicMock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=204)
kernel_id = token_hex(12)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel(kernel_id).destroy()
mock_req_cls.assert_called_once_with(
session,
'DELETE',
'/kernel/{}'.format(kernel_id),
params={})
async def test_get_kernel_info_url(self, mocker):
return_value = {}
mock_json_coro = asynctest.CoroutineMock(return_value=return_value)
mock_req_obj = asynctest.MagicMock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=200,
json=mock_json_coro)
kernel_id = token_hex(12)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel(kernel_id).get_info()
mock_req_cls.assert_called_once_with(
session, 'GET', '/kernel/{}'.format(kernel_id), params={})
async def test_restart_kernel_url(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=204)
session_id = secrets.token_hex(12)
mock_req_cls = mocker.patch('ai.backend.client.func.session.Request',
return_value=mock_req_obj)
async with AsyncSession() as session:
prefix = get_naming(session.api_version, 'path')
await session.ComputeSession(session_id).restart()
mock_req_cls.assert_called_once_with(session,
'PATCH',
f'/{prefix}/{session_id}',
params={})
async def test_create_kernel_return_id_only():
mock_resp = asynctest.MagicMock(spec=Response)
mock_resp.status = 201
mock_resp.json = lambda: {'kernelId': 'mock_kernel_id'}
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.afetch.return_value = mock_resp
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj):
k = await session.Kernel.get_or_create('python')
assert k.kernel_id == mock_resp.json()['kernelId']
async def test_stream_pty(mocker):
mock_req_obj = asynctest.MagicMock(spec=Request)
ws = object()
mock_req_obj.connect_websocket.return_value = None, ws
kernel_id = token_hex(12)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
stream = await session.Kernel(kernel_id).stream_pty()
mock_req_cls.assert_called_once_with(
session, 'GET', '/stream/kernel/{}/pty'.format(kernel_id))
mock_req_obj.connect_websocket.assert_called_once_with()
assert isinstance(stream, StreamPty)
assert stream.ws is ws
async def test_execute_code_url(self, mocker):
return_value = {'result': 'hi'}
mock_json_coro = asynctest.CoroutineMock(return_value=return_value)
mock_req_obj = asynctest.MagicMock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=200,
json=mock_json_coro)
kernel_id = token_hex(12)
run_id = token_hex(8)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel(kernel_id).execute(run_id, 'hello')
mock_req_cls.assert_called_once_with(
session, 'POST', '/kernel/{}'.format(kernel_id), params={})
async def test_create_kernel_url(self, mocker):
mock_req_obj = asynctest.MagicMock()
mock_req_obj.fetch.return_value = ContextMagicMock(
status=201, json=asynctest.CoroutineMock())
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel.get_or_create('python:3.6-ubuntu18.04')
mock_req_cls.assert_called_once_with(session, 'POST',
'/kernel/create')
mock_req_obj.fetch.assert_called_once_with()
mock_req_obj.fetch.return_value.json.assert_called_once_with()
async def test_create_kernel_url(mocker):
mock_resp = asynctest.MagicMock(spec=Response)
mock_resp.status = 201
mock_resp.json = asynctest.MagicMock()
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.afetch.return_value = mock_resp
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel.get_or_create('python')
mock_req_cls.assert_called_once_with(
session, 'POST', '/kernel/create', mock.ANY)
mock_req_obj.afetch.assert_called_once_with()
mock_req_obj.afetch.return_value.json.assert_called_once_with()
async def test_execute_code_url(mocker):
return_value = {'result': 'hi'}
mock_json_coro = AsyncMock(return_value=return_value)
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=200,
json=mock_json_coro)
session_id = secrets.token_hex(12)
run_id = secrets.token_hex(8)
mock_req_cls = mocker.patch('ai.backend.client.func.session.Request',
return_value=mock_req_obj)
async with AsyncSession() as session:
prefix = get_naming(session.api_version, 'path')
await session.ComputeSession(session_id).execute(run_id, 'hello')
mock_req_cls.assert_called_once_with(session,
'POST',
f'/{prefix}/{session_id}',
params={})
async def test_upload_jwt_generation(tmp_path):
with aioresponses() as m:
async with AsyncSession() as session:
mock_file = tmp_path / 'example.bin'
mock_file.write_bytes(secrets.token_bytes(32))
vfolder_name = 'fake-vfolder-name'
file_size = '1024'
payload = {
'token':
'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9. \
eyJwYXRoIjoiaHR0cDoxMjcuMC4wLjEvZm9sZGVycy9mYWtlLXZmb2xkZXItbmFtZS9yZXF1ZXN0LXVwbG9hZCIsInNpemUiOjEwMjR9.\
5IXk0xdrr6aPzVjud4cdfcXWch7Bq-m7SlFhnUv8XL8'
}
m.post(
build_url(session.config,
'/folders/{}/request-upload'.format(vfolder_name)),
payload=payload,
status=200)
rqst = Request('POST',
'/folders/{}/request-upload'.format(vfolder_name))
rqst.set_json({
'path': "{}".format(str(Path(mock_file))),
'size': str(file_size),
})
async with rqst.fetch() as resp:
res = await resp.json()
assert isinstance(resp, Response)
assert resp.status == 200
assert resp.content_type == 'application/json'
assert res == payload
assert 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9' in res['token']
async def test_vfolder_download(mocker):
mock_reader = AsyncMock()
mock_from_response = mocker.patch(
'ai.backend.client.func.vfolder.aiohttp.MultipartReader.from_response',
return_value=mock_reader)
mock_reader.next = AsyncMock()
mock_reader.next.return_value = None
mock_file = 'fake-file1'
with aioresponses() as m:
async with AsyncSession() as session:
vfolder_name = 'fake-vfolder-name'
# client to manager
# manager to storage-proxy
storage_path = str(
build_url(session.config,
'folder/{}/download'.format(vfolder_name))).replace(
'8081', '6021')
storage_path2 = str(build_url(session.config,
'/download')).replace(
'8081', '6021')
payload = {
'token': 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9. \
eyJwYXRoIjoiaHR0cDoxMjcuMC4wLjEvZm9sZGVycy9mYWtlLXZmb2xkZXItbmFtZS9yZXF1ZXN0LXVwbG9hZCIsInNpemUiOjEwMjR9.\
5IXk0xdrr6aPzVjud4cdfcXWch7Bq-m7SlFhnUv8XL8',
'url': storage_path
}
storage_payload = {
'token':
'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9. \
eyJvcCI6InVwbG9hZCIsInZvbHVtZSI6InZvbHVtZTEiLCJ2ZmlkIjoiO \
DBiYWYyYjgtNTY3My00MmVkLTgyZWEtYj \
NmNzNmOWQwNjAzIiwicmVscGF0aCI6InNldHVwLmNmZyIsInNpemUiOjU \
yNywic2Vzc2lvbiI6ImE3YzZiY2I1MWRlY2I3NzJjZjRkMDI3YjA5 \
MGI5NGM5IiwiZXhwIjoxNTk5MTIzMzYxfQ. \
D13UMFrz-2qq9c0k4MGpjVOMn5Z9-fyR5tRRIkvtvqk'
}
# 1. Client to Manager throught Request
m.post(build_url(
session.config,
"/folders/{}/request-download?path='{}'".format(
vfolder_name, mock_file)),
payload=payload['token'],
status=200)
# 2. Manager to storage proxy
"""
m.post(storage_path + "?volume= \
volume1&vfid=80baf2b8-5673-42ed-82ea-b3f73f9d0603&relpath={}"
.format('fake-file1'),
payload=payload,
status=200)
"""
# 3. Client to Manager through TusClient. Upload url
m.get(storage_path2 + "?token={}".format(storage_payload['token']))
m.get(build_url(
session.config,
"/folders/{}/request-download?path='{}'".format(
vfolder_name, mock_file)),
status=200)
await session.VFolder(vfolder_name).download(['fake-file1'])
assert mock_from_response.called == 1
assert mock_reader.next.called == 1
async def test_vfolder_upload(tmp_path: Path):
mock_file = tmp_path / 'example.bin'
mock_file.write_bytes(secrets.token_bytes(1024))
with aioresponses() as m:
async with AsyncSession() as session:
vfolder_name = 'fake-vfolder-name'
storage_path = str(
build_url(session.config,
'folder/{}/upload'.format(vfolder_name))).replace(
'8081', '6021')
storage_path2 = str(build_url(session.config,
'/upload')).replace('8081', '6021')
payload = {
'token': 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9. \
eyJwYXRoIjoiaHR0cDoxMjcuMC4wLjEvZm9sZGVycy9mYWtlLXZmb2xkZXItbmFtZS9yZXF1ZXN0LXVwbG9hZCIsInNpemUiOjEwMjR9.\
5IXk0xdrr6aPzVjud4cdfcXWch7Bq-m7SlFhnUv8XL8',
'url': storage_path
}
storage_payload = {
'token':
'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9. \
eyJvcCI6InVwbG9hZCIsInZvbHVtZSI6InZvbHVtZTEiLCJ2ZmlkIjoiO \
DBiYWYyYjgtNTY3My00MmVkLTgyZWEtYj \
NmNzNmOWQwNjAzIiwicmVscGF0aCI6InNldHVwLmNmZyIsInNpemUiOjU \
yNywic2Vzc2lvbiI6ImE3YzZiY2I1MWRlY2I3NzJjZjRkMDI3YjA5 \
MGI5NGM5IiwiZXhwIjoxNTk5MTIzMzYxfQ. \
D13UMFrz-2qq9c0k4MGpjVOMn5Z9-fyR5tRRIkvtvqk'
}
"""
# 0. This works and passes test when reqeusting jwt in test_upload_jwt_generation().
# but here it freezes the client
"""
m.post(
build_url(session.config,
'/folders/{}/request-upload'.format(vfolder_name)),
payload=payload,
status=200)
# 1. Client to Manager throught Request
m.post(build_url(
session.config,
"/folders/{}/request-upload?path='{}'&size={}".format(
vfolder_name, mock_file, 1024)),
payload=payload['token'],
status=200)
# 2. Response from storage to manager
m.post(storage_path + "?volume= \
volume1&vfid=80baf2b8-5673-42ed-82ea-b3f73f9d0603&relpath={}&size={}"
.format('example.bin', '1024'),
payload=payload,
status=200)
# 3. Client to Manager through TusClient. Upload url
tus_payload = {
'Access-Control-Allow-Origin': '*',
'Access-Control-Allow-Headers': 'Tus-Resumable, \
Upload-Length, Upload-Metadata, Upload-Offset, Content-Type',
'Access-Control-Expose-Headers':
'Tus-Resumable, Upload-Length, Upload-Metadata, \
Upload-Offset, Content-Type',
'Access-Control-Allow-Methods': '*',
'Cache-Control': 'no-store',
'Tus-Resumable': '1.0.0',
'Upload-Offset': '527',
'Upload-Length': '527',
'Content-Length': '0',
'Content-Type': 'application/octet-stream',
'Date': 'Wed, 02 Sep 2020 12:54:17 GMT',
'Server': 'Python/3.8 aiohttp/3.6.2'
}
# file upload on storage
m.patch(storage_path2 +
"?token={}".format(storage_payload['token']),
payload=tus_payload,
headers=tus_payload,
status=204)
m.patch(build_url(
session.config,
"/folders/{}/request-upload?path='{}'&size={}".format(
vfolder_name, mock_file, 1024)),
status=200)
resp = await session.VFolder(vfolder_name).upload([mock_file],
basedir=tmp_path)
assert resp == ""
async def test_async_connection(self):
async with AsyncSession() as sess:
request = Request(sess, 'GET', '/')
async with request.fetch() as resp:
assert 'version' in await resp.json()
async def test_fetch_invalid_method_async():
async with AsyncSession() as session:
rqst = Request('STRANGE', '/')
with pytest.raises(AssertionError):
async with rqst.fetch():
pass
