async def test_fetch(dummy_endpoint): with aioresponses() as m, Session() as session: body = b'hello world' m.post( dummy_endpoint + 'function', status=200, body=body, headers={'Content-Type': 'text/plain; charset=utf-8', 'Content-Length': str(len(body))}, ) rqst = Request('POST', 'function') async with rqst.fetch() as resp: assert isinstance(resp, Response) assert resp.status == 200 assert resp.content_type == 'text/plain' assert await resp.text() == body.decode() assert resp.content_length == len(body) with aioresponses() as m, Session() as session: body = b'{"a": 1234, "b": null}' m.post( dummy_endpoint + 'function', status=200, body=body, headers={'Content-Type': 'application/json; charset=utf-8', 'Content-Length': str(len(body))}, ) rqst = Request('POST', 'function') async with rqst.fetch() as resp: assert isinstance(resp, Response) assert resp.status == 200 assert resp.content_type == 'application/json' assert await resp.text() == body.decode() assert await resp.json() == {'a': 1234, 'b': None} assert resp.content_length == len(body)
def test_api_function_metaclass(): # Here, we repeat intentionally the same stuffs # to check if our metaclass works across multiple # re-definition and re-instantiation scenarios. with Session() as session: Dummy = type('DummyFunction', (BaseFunction, ), { **DummyFunction.__dict__, 'session': session, }) assert Dummy.session is session assert Dummy().session is session assert Dummy.get_or_create() == 'created' assert Dummy().calculate() == 'done' assert Dummy.get_or_create() == 'created' assert Dummy().calculate() == 'done' with Session() as session: Dummy = type('DummyFunction', (BaseFunction, ), { **DummyFunction.__dict__, 'session': session, }) assert Dummy.session is session assert Dummy().session is session assert Dummy.get_or_create() == 'created' assert Dummy().calculate() == 'done' assert Dummy.get_or_create() == 'created' assert Dummy().calculate() == 'done'
def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.log.info( 'Backend.AI kernel starting with client session ID: {0}'.format( self.ident)) self.backend_session = Session() self.kernel = self.backend_session.Kernel.get_or_create( self.backend_lang, self.ident)
def test_vfolder_list_files(): with Session() as session: with aioresponses() as m: vfolder_name = 'fake-vfolder-name' payload = { "files": [{ "mode": "-rw-r--r--", "size": 4751244, "ctime": 1528277299.2744732, "mtime": 1528277299.2744732, "atime": 1528277300.7658687, "filename": "bigtxt.txt", }, { "mode": "-rw-r--r--", "size": 200000, "ctime": 1528333257.6576185, "mtime": 1528288069.625786, "atime": 1528332829.692922, "filename": "200000", }], "folder_path": "/mnt/local/1f6bd27fde1248cabfb50306ea83fc0a", } m.get(build_url(session.config, '/folders/{}/files'.format(vfolder_name)), status=200, payload=payload) resp = session.VFolder(vfolder_name).list_files('.') assert resp == payload
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_delete_vfolder(): with Session() as session, aioresponses() as m: vfolder_name = 'fake-vfolder-name' m.delete(build_url(session.config, '/folders/{}'.format(vfolder_name)), status=204) resp = session.VFolder(vfolder_name).delete() assert resp == {}
def test_kernel_execution_with_vfolder_mounts(self): with Session() as sess: vfname = 'vftest-' + token_hex(4) sess.VFolder.create(vfname) vfolder = sess.VFolder(vfname) try: with tempfile.NamedTemporaryFile('w', suffix='.py', dir=Path.cwd()) as f: f.write('print("hello world")\nx = 1 / 0\n') f.flush() f.seek(0) vfolder.upload([f.name]) kernel = sess.Kernel.get_or_create('python:3.6-ubuntu18.04', mounts=[vfname]) try: console, n = exec_loop( kernel, 'batch', '', { 'build': '', 'exec': 'python {}/{}'.format(vfname, Path(f.name).name), }) assert 'hello world' in console['stdout'] assert 'ZeroDivisionError' in console['stderr'] assert len(console['media']) == 0 finally: kernel.destroy() finally: vfolder.delete()
def test_auth_missing_body(): with Session() as sess: request = Request('GET', '/auth') with pytest.raises(BackendAPIError) as e: with request.fetch(): pass assert e.value.status == 400
async def test_user_cannot_mutate_alias_dealias(self, userconfig): with Session() as sess: test_alias = 'testalias-b9f1ce136f584ca892d5fef3e78dd11d' with pytest.raises(BackendAPIError): sess.Image.aliasImage(test_alias, 'lua:5.1-alpine3.8') with pytest.raises(BackendAPIError): sess.Image.dealiasImage(test_alias)
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_raises_err_with_abnormal_status(mocker): mock_req_obj = mock.Mock() mock_req_obj.fetch.return_value = mock.MagicMock(status=400) mocker.patch('ai.backend.client.kernel.Request', return_value=mock_req_obj) with Session() as session: with pytest.raises(BackendAPIError): session.Kernel.get_or_create('python')
async def test_list_images_by_user(self, userconfig): with Session() as sess: images = sess.Image.list() image = images[0] assert len(images) > 0 assert 'name' in image assert 'tag' in image assert 'hash' in image
def test_vfolder_delete_invitation(): with Session() as session, aioresponses() as m: payload = {'msg': 'Vfolder invitation is deleted: fake-inv-id.'} m.delete(build_url(session.config, '/folders/invitations/delete'), status=200, payload=payload) resp = session.VFolder.delete_invitation('inv-id') assert resp == payload
async def test_manipulate_resource_policy(self): access_key = get_config().access_key rpname = 'testrp-' + uuid.uuid4().hex with Session() as sess: try: rp = sess.ResourcePolicy(access_key) assert rp.info(rpname) is None rps = sess.ResourcePolicy.list() original_count = len(rps) # Create resource policy sess.ResourcePolicy.create(name=rpname, default_for_unspecified='LIMITED', total_resource_slots='{}', max_concurrent_sessions=1, max_containers_per_session=1, max_vfolder_count=1, max_vfolder_size=1, idle_timeout=1, allowed_vfolder_hosts=['local']) rps = sess.ResourcePolicy.list() assert len(rps) == original_count + 1 info = rp.info(rpname) assert info['name'] == rpname assert info['total_resource_slots'] == '{}' assert info['max_concurrent_sessions'] == 1 assert info['max_vfolder_count'] == 1 assert info['max_vfolder_size'] == 1 assert info['idle_timeout'] == 1 # Update resource policy sess.ResourcePolicy.update( name=rpname, default_for_unspecified='LIMITED', total_resource_slots='{"cpu": "count"}', max_concurrent_sessions=2, max_containers_per_session=2, max_vfolder_count=2, max_vfolder_size=2, idle_timeout=2, allowed_vfolder_hosts=['local']) rps = sess.ResourcePolicy.list() assert len(rps) == original_count + 1 info = rp.info(rpname) assert info['name'] == rpname assert info['total_resource_slots'] == '{"cpu": "count"}' assert info['max_concurrent_sessions'] == 2 assert info['max_vfolder_count'] == 2 assert info['max_vfolder_size'] == 2 assert info['idle_timeout'] == 2 # Delete ResourcePolicy sess.ResourcePolicy.delete(rpname) rps = sess.ResourcePolicy.list() assert len(rps) == original_count except Exception: sess.ResourcePolicy.delete(rpname) raise
async def test_user_cannot_create_keypair(self, userconfig): email = 'testion' + uuid.uuid4().hex + '@test.mars' with Session() as sess: with pytest.raises(BackendAPIError): sess.KeyPair.create(user_id=email, is_active=True, is_admin=False, resource_policy='default', rate_limit=1)
def test_freeze_opt_force_kill(self, mocker): mock_req_obj = mock.Mock() mock_req_obj.fetch.return_value = ContextMagicMock(status=204) mocker.patch('ai.backend.client.manager.Request', return_value=mock_req_obj) with Session() as session: session.Manager.freeze(force_kill=True) mock_req_obj.fetch.assert_called_once_with()
def test_destroy_kernel_raises_err_with_abnormal_status(mocker): mock_req_obj = mock.Mock() mock_req_obj.fetch.return_value = mock.MagicMock(status=400) mocker.patch('ai.backend.client.kernel.Request', return_value=mock_req_obj) kernel_id = token_hex(12) with Session() as session: with pytest.raises(BackendAPIError): k = session.Kernel(kernel_id) k.destroy()
def test_vfolder_delete_files(): with Session() as session, aioresponses() as m: vfolder_name = 'fake-vfolder-name' files = ['fake-file1', 'fake-file2'] m.delete(build_url(session.config, '/folders/{}/delete_files'.format(vfolder_name)), status=200, payload={}) resp = session.VFolder(vfolder_name).delete_files(files) assert resp == '{}'
def test_create_kernel_return_id_only(mocker): mock_json_func = lambda: {'kernelId': 'mock_kernel_id'} mock_req_obj = mock.Mock() mock_req_obj.fetch.return_value = mock.MagicMock(status=201, json=mock_json_func) mocker.patch('ai.backend.client.kernel.Request', return_value=mock_req_obj) with Session() as session: k = session.Kernel.get_or_create('python') assert k.kernel_id == mock_json_func()['kernelId']
def test_vfolder_upload(tmp_path: Path): with Session() as session, aioresponses() as m: mock_file = tmp_path / 'example.bin' mock_file.write_bytes(secrets.token_bytes(32)) vfolder_name = 'fake-vfolder-name' m.post(build_url(session.config, '/folders/{}/upload'.format(vfolder_name)), status=201) resp = session.VFolder(vfolder_name).upload([mock_file], basedir=tmp_path) assert resp == ''
def test_vfolder_invite(): with Session() as session, aioresponses() as m: vfolder_name = 'fake-vfolder-name' user_ids = ['*****@*****.**', '*****@*****.**'] payload = {'invited_ids': user_ids} m.post(build_url(session.config, '/folders/{}/invite'.format(vfolder_name)), status=201, payload=payload) resp = session.VFolder(vfolder_name).invite('rw', user_ids) assert resp == payload
def test_auth_malformed(): with Session() as sess: request = Request('GET', '/auth') request.set_content( b'<this is not json>', content_type='application/json', ) with pytest.raises(BackendAPIError) as e: with request.fetch(): pass assert e.value.status == 400
def test_vfolder_accept_invitation(): with Session() as session, aioresponses() as m: payload = { 'msg': ('User [email protected] now can access' ' vfolder fake-vfolder-id'), } m.post(build_url(session.config, '/folders/invitations/accept'), status=200, payload=payload) resp = session.VFolder.accept_invitation('inv-id') assert resp == payload
def test_vfolder_upload(tmpdir): with Session() as session: with aioresponses() as m: mockfile = tmpdir.join('example.jpg') mockfile.write('mock file') vfolder_name = 'fake-vfolder-name' m.post(build_url(session.config, '/folders/{}/upload'.format(vfolder_name)), status=201) resp = session.VFolder(vfolder_name).upload([mockfile.strpath], basedir=tmpdir.strpath) assert resp.status == 201
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_kernel_get_or_create_reuse(self): with Session() as sess: try: # Sessions with same token and same language must be reused. t = token_hex(6) kernel1 = sess.Kernel.get_or_create('python:3.6-ubuntu18.04', client_token=t) kernel2 = sess.Kernel.get_or_create('python:3.6-ubuntu18.04', client_token=t) assert kernel1.kernel_id == kernel2.kernel_id finally: kernel1.destroy()
def test_auth(): random_msg = uuid.uuid4().hex with Session() as sess: request = Request('GET', '/auth') request.set_json({ 'echo': random_msg, }) with request.fetch() as resp: assert resp.status == 200 data = resp.json() assert data['authorized'] == 'yes' assert data['echo'] == random_msg
def test_create_kernel_url(self, mocker): mock_req_obj = mock.Mock() mock_req_obj.fetch.return_value = ContextMagicMock( status=201, json=asynctest.CoroutineMock()) mock_req = mocker.patch('ai.backend.client.kernel.Request', return_value=mock_req_obj) with Session() as session: session.Kernel.get_or_create('python:3.6-ubuntu18.04') mock_req.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()
def test_not_found(self): with Session() as sess: request = Request(sess, 'GET', '/invalid-url-wow') with pytest.raises(BackendAPIError) as e: with request.fetch(): pass assert e.value.status == 404 request = Request(sess, 'GET', '/auth/uh-oh') with pytest.raises(BackendAPIError) as e: with request.fetch(): pass assert e.value.status == 404
def test_create_kernel_return_id_only(self, mocker): return_value = {'kernelId': 'mock_kernel_id'} mock_json_coro = asynctest.CoroutineMock(return_value=return_value) mock_req_obj = mock.Mock() mock_req_obj.fetch.return_value = ContextMagicMock(status=201, json=mock_json_coro) mocker.patch('ai.backend.client.kernel.Request', return_value=mock_req_obj) with Session() as session: k = session.Kernel.get_or_create('python:3.6-ubuntu18.04') assert k.kernel_id == return_value['kernelId']