def test_fetch(self, 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(session, 'POST', 'function') with rqst.fetch() as resp: assert isinstance(resp, Response) assert resp.status == 200 assert resp.content_type == 'text/plain' assert 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(session, 'POST', 'function') with rqst.fetch() as resp: assert isinstance(resp, Response) assert resp.status == 200 assert resp.content_type == 'application/json' assert resp.text() == body.decode() assert resp.json() == {'a': 1234, 'b': None} assert resp.content_length == len(body)
async def info(self, fields: Iterable[str] = None) -> dict: ''' Returns the keypair's information such as resource limits. :param fields: Additional per-agent query fields to fetch. .. versionadded:: 18.12 ''' if fields is None: fields = ( 'access_key', 'secret_key', 'is_active', 'is_admin', ) q = 'query {' \ ' keypair {' \ ' $fields' \ ' }' \ '}' q = q.replace('$fields', ' '.join(fields)) rqst = Request(self.session, 'POST', '/admin/graphql') rqst.set_json({ 'query': q, }) async with rqst.fetch() as resp: data = await resp.json() return data['keypair']
async def update(cls, access_key: str, is_active: bool = None, is_admin: bool = None, resource_policy: str = None, rate_limit: int = None) -> dict: """ Creates a new keypair with the given options. You need an admin privilege for this operation. """ q = 'mutation($access_key: String!, $input: ModifyKeyPairInput!) {' + \ ' modify_keypair(access_key: $access_key, props: $input) {' \ ' ok msg' \ ' }' \ '}' variables = { 'access_key': access_key, 'input': { 'is_active': is_active, 'is_admin': is_admin, 'resource_policy': resource_policy, 'rate_limit': rate_limit, }, } rqst = Request(cls.session, 'POST', '/admin/graphql') rqst.set_json({ 'query': q, 'variables': variables, }) async with rqst.fetch() as resp: data = await resp.json() return data['modify_keypair']
def test_fetch_invalid_method(self, mock_request_params): mock_request_params['method'] = 'STRANGE' rqst = Request(**mock_request_params) with pytest.raises(AssertionError): with rqst.fetch(): pass
async def list(cls, operation: bool = False, fields: Iterable[str] = None) -> Sequence[dict]: ''' Fetches the list of registered images in this cluster. ''' if fields is None: fields = ( 'name', 'tag', 'hash', ) q = 'query($is_operation: Boolean) {' \ ' images(is_operation: $is_operation) {' \ ' $fields' \ ' }' \ '}' q = q.replace('$fields', ' '.join(fields)) variables = { 'is_operation': operation, } rqst = Request(cls.session, 'POST', '/admin/graphql') rqst.set_json({ 'query': q, 'variables': variables, }) async with rqst.fetch() as resp: data = await resp.json() return data['images']
async def test_fetch_invalid_method(mock_request_params): mock_request_params['method'] = 'STRANGE' rqst = Request(**mock_request_params) with pytest.raises(AssertionError): async with rqst.fetch(): pass
async def deactivate(cls, access_key: str) -> dict: ''' Deactivates this keypair. Deactivated keypairs cannot make any API requests unless activated again by an administrator. You need an admin privilege for this operation. ''' q = 'mutation($access_key: String!, $input: ModifyKeyPairInput!) {' + \ ' modify_keypair(access_key: $access_key, props: $input) {' \ ' ok msg' \ ' }' \ '}' variables = { 'access_key': access_key, 'input': { 'is_active': False, 'is_admin': None, 'resource_policy': None, 'rate_limit': None, }, } rqst = Request(cls.session, 'POST', '/admin/graphql') rqst.set_json({ 'query': q, 'variables': variables, }) async with rqst.fetch() as resp: data = await resp.json() return data['modify_keypair']
def test_auth_missing_signature(defconfig): random_msg = uuid.uuid4().hex request = Request('GET', '/authorize', { 'echo': random_msg, }) resp = request.send() assert resp.status == 401
async def detail(cls, gid: str, fields: Iterable[str] = None) -> Sequence[dict]: ''' Fetch information of a group with group ID. :param gid: ID of the group to fetch. :param fields: Additional per-group query fields to fetch. ''' if fields is None: fields = ('id', 'name', 'description', 'is_active', 'created_at', 'domain_name', 'total_resource_slots', 'allowed_vfolder_hosts', 'integration_id') query = textwrap.dedent('''\ query($gid: String!) { group(id: $gid) {$fields} } ''') query = query.replace('$fields', ' '.join(fields)) variables = {'gid': gid} rqst = Request(cls.session, 'POST', '/admin/graphql') rqst.set_json({ 'query': query, 'variables': variables, }) async with rqst.fetch() as resp: data = await resp.json() return data['group']
async def list(cls, domain_name: str, fields: Iterable[str] = None) -> Sequence[dict]: ''' Fetches the list of groups. :param domain_name: Name of domain to list groups. :param fields: Additional per-group query fields to fetch. ''' if fields is None: fields = ('id', 'name', 'description', 'is_active', 'created_at', 'domain_name', 'total_resource_slots', 'allowed_vfolder_hosts', 'integration_id') query = textwrap.dedent('''\ query($domain_name: String) { groups(domain_name: $domain_name) {$fields} } ''') query = query.replace('$fields', ' '.join(fields)) variables = {'domain_name': domain_name} rqst = Request(cls.session, 'POST', '/admin/graphql') rqst.set_json({ 'query': query, 'variables': variables, }) async with rqst.fetch() as resp: data = await resp.json() return data['groups']
async def remove_users(cls, gid: str, user_uuids: Iterable[str], fields: Iterable[str] = None) -> dict: ''' Remove users from a group. You need an admin privilege for this operation. ''' query = textwrap.dedent('''\ mutation($gid: String!, $input: ModifyGroupInput!) { modify_group(gid: $gid, props: $input) { ok msg } } ''') variables = { 'gid': gid, 'input': { 'user_update_mode': 'remove', 'user_uuids': user_uuids, }, } rqst = Request(cls.session, 'POST', '/admin/graphql') rqst.set_json({ 'query': query, 'variables': variables, }) async with rqst.fetch() as resp: data = await resp.json() return data['modify_group']
async def update(cls, gid: str, name: str = None, description: str = None, is_active: bool = None, total_resource_slots: str = None, allowed_vfolder_hosts: Iterable[str] = None, integration_id: str = None, fields: Iterable[str] = None) -> dict: ''' Update existing group. You need an admin privilege for this operation. ''' query = textwrap.dedent('''\ mutation($gid: String!, $input: ModifyGroupInput!) { modify_group(gid: $gid, props: $input) { ok msg } } ''') variables = { 'gid': gid, 'input': { 'name': name, 'description': description, 'is_active': is_active, 'total_resource_slots': total_resource_slots, 'allowed_vfolder_hosts': allowed_vfolder_hosts, 'integration_id': integration_id, }, } rqst = Request(cls.session, 'POST', '/admin/graphql') rqst.set_json({ 'query': query, 'variables': variables, }) async with rqst.fetch() as resp: data = await resp.json() return data['modify_group']
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
def test_build_correct_url(req_params): config = req_params['config'] req = Request(**req_params) major_ver = config.version.split('.', 1)[0] path = '/' + req.path if len(req.path) > 0 else '' assert req.build_url() == urljoin(config.endpoint, major_ver + path)
def test_set_content_correctly(req_params): req_params['content'] = OrderedDict() req = Request(**req_params) new_data = b'new-data' assert not req.content req.content = new_data assert req.content == new_data assert req.headers['Content-Length'] == str(len(new_data))
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_fetch_timeout_async(self, 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
def test_auth_missing_signature(defconfig): random_msg = uuid.uuid4().hex request = Request('GET', '/authorize', { 'echo': random_msg, }) # let it bypass actual signing request._sign = lambda *args, **kwargs: None resp = request.send() assert resp.status == 401
async def test_fetch_client_error_async(self, 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
def test_auth(defconfig): random_msg = uuid.uuid4().hex request = Request('GET', '/authorize', { 'echo': random_msg, }) resp = request.send() assert resp.status == 200 data = resp.json() assert data['authorized'] == 'yes' assert data['echo'] == random_msg
def test_build_correct_url(self, mock_request_params): canonical_url = 'http://127.0.0.1:8081/function?app=999' mock_request_params['path'] = '/function' rqst = Request(**mock_request_params) assert str(rqst._build_url()) == canonical_url mock_request_params['path'] = 'function' rqst = Request(**mock_request_params) assert str(rqst._build_url()) == canonical_url
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
def test_auth(self): random_msg = uuid.uuid4().hex with Session() as sess: request = Request(sess, '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_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_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_auth_missing_signature(monkeypatch): random_msg = uuid.uuid4().hex with Session() as sess: rqst = Request('GET', '/auth') rqst.set_json({'echo': random_msg}) # let it bypass actual signing from ai.backend.client import request noop_sign = lambda *args, **kwargs: ({}, None) monkeypatch.setattr(request, 'generate_signature', noop_sign) with pytest.raises(BackendAPIError) as e: with rqst.fetch(): pass assert e.value.status == 401
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_response_async(self, 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(session, 'POST', '/function') async with rqst.fetch() as resp: assert await resp.text() == '{"test": 5678}' assert await resp.json() == {'test': 5678}
def test_auth_missing_signature(self, monkeypatch): random_msg = uuid.uuid4().hex with Session() as sess: rqst = Request(sess, 'GET', '/auth') rqst.set_json({'echo': random_msg}) # let it bypass actual signing from ai.backend.client import request noop_sign = lambda *args, **kwargs: ({}, None) monkeypatch.setattr(request, 'generate_signature', noop_sign) with pytest.raises(BackendAPIError) as e: with rqst.fetch(): pass assert e.value.status == 401
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
def test_content_is_auto_set_to_blank_if_no_data(req_params): req_params = req_params.copy() req_params['content'] = None req = Request(**req_params) assert req.content_type == 'application/octet-stream' assert req.content == b''
async def rescan_images(cls, registry: str): q = 'mutation($registry: String) {' \ ' rescan_images(registry:$registry) {' \ ' ok msg' \ ' }' \ '}' variables = { 'registry': registry, } rqst = Request(cls.session, 'POST', '/admin/graphql') rqst.set_json({ 'query': q, 'variables': variables, }) async with rqst.fetch() as resp: data = await resp.json() return data['rescan_images']
async def dealias_image(cls, alias: str) -> dict: q = 'mutation($alias: String!) {' \ ' dealias_image(alias: $alias) {' \ ' ok msg' \ ' }' \ '}' variables = { 'alias': alias, } rqst = Request(cls.session, 'POST', '/admin/graphql') rqst.set_json({ 'query': q, 'variables': variables, }) async with rqst.fetch() as resp: data = await resp.json() return data['dealias_image']
def test_request_initialization(mock_request_params): rqst = Request(**mock_request_params) assert rqst.method == mock_request_params['method'] assert rqst.params == mock_request_params['params'] assert rqst.path == mock_request_params['path'].lstrip('/') assert rqst.content == mock_request_params['content'] assert 'X-BackendAI-Version' in rqst.headers
async def test_streaming_fetch(dummy_endpoint): # Read content by chunks. 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 resp.status == 200 assert resp.content_type == 'text/plain' assert await resp.read(3) == b'hel' assert await resp.read(2) == b'lo' await resp.read() with pytest.raises(AssertionError): assert await resp.text()
def test_streaming_fetch(self, dummy_endpoint): # Read content by chunks. 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(session, 'POST', 'function') with rqst.fetch() as resp: assert resp.status == 200 assert resp.content_type == 'text/plain' assert resp.read(3) == b'hel' assert resp.read(2) == b'lo' resp.read() with pytest.raises(AssertionError): assert resp.text()
async def test_invalid_requests(dummy_endpoint): with aioresponses() as m, Session() as session: body = json.dumps({ 'type': 'https://api.backend.ai/probs/kernel-not-found', 'title': 'Kernel Not Found', }).encode('utf8') m.post( dummy_endpoint, status=404, body=body, headers={'Content-Type': 'application/problem+json; charset=utf-8', 'Content-Length': str(len(body))}, ) rqst = Request('POST', '/') with pytest.raises(BackendAPIError) as e: async with rqst.fetch(): pass assert e.status == 404 assert e.data['type'] == \ 'https://api.backend.ai/probs/kernel-not-found' assert e.data['title'] == 'Kernel Not Found'
def test_content_is_files(req_params): files = [ ('src', 'test1.txt', io.BytesIO(), 'application/octet-stream'), ('src', 'test2.txt', io.BytesIO(), 'application/octet-stream'), ] req_params['content'] = files req = Request(**req_params) assert req.content_type == 'multipart/form-data' assert req.content == files
def test_invalid_requests(self, dummy_endpoint): with aioresponses() as m, Session() as session: body = json.dumps({ 'type': 'https://api.backend.ai/probs/kernel-not-found', 'title': 'Kernel Not Found', }).encode('utf8') m.post( dummy_endpoint, status=404, body=body, headers={'Content-Type': 'application/problem+json; charset=utf-8', 'Content-Length': str(len(body))}, ) rqst = Request(session, 'POST', '/') with pytest.raises(BackendAPIError) as e: with rqst.fetch(): pass assert e.status == 404 assert e.data['type'] == \ 'https://api.backend.ai/probs/kernel-not-found' assert e.data['title'] == 'Kernel Not Found'
def test_request_initialization(req_params): req = Request(**req_params) assert req.config == req_params['config'] assert req.method == req_params['method'] assert req.path == req_params['path'][1:] assert req.content == req_params['content'] assert 'Date' in req.headers assert 'X-BackendAI-Version' in req.headers assert req._content == json.dumps(req_params['content']).encode('utf8')
def test_request_attach_files(self, mock_request_params): files = [ AttachedFile('test1.txt', io.BytesIO(), 'application/octet-stream'), AttachedFile('test2.txt', io.BytesIO(), 'application/octet-stream'), ] mock_request_params['content'] = b'something' rqst = Request(**mock_request_params) with pytest.raises(AssertionError): rqst.attach_files(files) mock_request_params['content'] = b'' rqst = Request(**mock_request_params) rqst.attach_files(files) assert rqst.content_type == 'multipart/form-data' assert rqst.content == b'' packed_content = rqst._pack_content() assert packed_content.is_multipart
def test_request_set_content(self, mock_request_params): rqst = Request(**mock_request_params) assert rqst.content == mock_request_params['content'] assert rqst.content_type == 'application/json' assert rqst._pack_content() is rqst.content mock_request_params['content'] = 'hello' mock_request_params['content_type'] = None rqst = Request(**mock_request_params) assert rqst.content == b'hello' assert rqst.content_type == 'text/plain' assert rqst._pack_content() is rqst.content mock_request_params['content'] = b'\x00\x01\xfe\xff' mock_request_params['content_type'] = None rqst = Request(**mock_request_params) assert rqst.content == b'\x00\x01\xfe\xff' assert rqst.content_type == 'application/octet-stream' assert rqst._pack_content() is rqst.content
async def test_fetch_invalid_method_async(self): async with AsyncSession() as session: rqst = Request(session, 'STRANGE', '/') with pytest.raises(AssertionError): async with rqst.fetch(): pass
def test_connection(self): with Session() as sess: request = Request(sess, 'GET', '/') with request.fetch() as resp: assert 'version' in resp.json()
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()
def test_request_set_content_none(self, mock_request_params): mock_request_params = mock_request_params.copy() mock_request_params['content'] = None rqst = Request(**mock_request_params) assert rqst.content == b'' assert rqst._pack_content() is rqst.content