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