def test_create_with_config(mocker, api_version):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=201,
json=AsyncMock())
mock_req = mocker.patch('ai.backend.client.func.session.Request',
return_value=mock_req_obj)
myconfig = APIConfig(
endpoint='https://localhost:9999',
access_key='1234',
secret_key='asdf',
user_agent='BAIClientTest',
version=f'v{api_version[0]}.{api_version[1]}',
)
with Session(config=myconfig) as session:
prefix = get_naming(session.api_version, 'path')
if api_version[0] == 4:
assert prefix == 'kernel'
else:
assert prefix == 'session'
assert session.config is myconfig
cs = session.ComputeSession.get_or_create('python')
mock_req.assert_called_once_with(session, 'POST', f'/{prefix}')
assert str(cs.session.config.endpoint) == 'https://localhost:9999'
assert cs.session.config.user_agent == 'BAIClientTest'
assert cs.session.config.access_key == '1234'
assert cs.session.config.secret_key == 'asdf'
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
with Session() as session, aioresponses() as m:
vfolder_name = 'fake-vfolder-name'
m.get(build_url(session.config,
'/folders/{}/download'.format(vfolder_name)),
status=200,
headers={'X-TOTAL-PAYLOADS-LENGTH': '0'},
body='')
session.VFolder(vfolder_name).download(['fake-file1'])
assert mock_from_response.called == 1
assert mock_reader.next.called == 1
def test_create_kernel_url(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=201,
json=AsyncMock())
mock_req = mocker.patch('ai.backend.client.func.session.Request',
return_value=mock_req_obj)
with Session() as session:
prefix = get_naming(session.api_version, 'path')
session.ComputeSession.get_or_create('python:3.6-ubuntu18.04')
mock_req.assert_called_once_with(session, 'POST', f'/{prefix}')
mock_req_obj.fetch.assert_called_once_with()
mock_req_obj.fetch.return_value.json.assert_called_once_with()
def test_status(mocker):
return_value = {'status': 'running', 'active_sessions': 3}
mock_json_coro = AsyncMock(return_value=return_value)
mock_req_obj = mocker.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=200,
json=mock_json_coro)
mocker.patch('ai.backend.client.func.manager.Request', return_value=mock_req_obj)
with Session() as session:
resp = session.Manager.status()
mock_req_obj.fetch.assert_called_once_with()
assert resp['status'] == return_value['status']
assert resp['active_sessions'] == return_value['active_sessions']
async def test_get_kernel_info_url(mocker):
return_value = {}
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)
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).get_info()
mock_req_cls.assert_called_once_with(session,
'GET',
f'/{prefix}/{session_id}',
params={})
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={})
def api_version(request):
mock_nego_func = AsyncMock()
mock_nego_func.return_value = request.param
with mock.patch('ai.backend.client.session._negotiate_api_version',
mock_nego_func):
yield request.param
def api_version():
mock_nego_func = AsyncMock()
mock_nego_func.return_value = API_VERSION
with mock.patch('ai.backend.client.session._negotiate_api_version', mock_nego_func):
yield
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