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