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_unfreeze(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=204)
mocker.patch('ai.backend.client.func.manager.Request', return_value=mock_req_obj)
with Session() as session:
session.Manager.unfreeze()
mock_req_obj.fetch.assert_called_once_with()
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_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={})
def test_destroy_kernel_url(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=204)
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_id = secrets.token_hex(12)
cs = session.ComputeSession(session_id)
cs.destroy()
mock_req.assert_called_once_with('DELETE',
f'/{prefix}/{session_id}',
params={})
mock_req_obj.fetch.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={})
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)
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_id = secrets.token_hex(12)
cs = session.ComputeSession(session_id)
cs.get_info()
mock_req.assert_called_once_with('GET',
f'/{prefix}/{session_id}',
params={})
mock_req_obj.fetch.assert_called_once_with()
mock_req_obj.fetch.return_value.json.assert_called_once_with()