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_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)
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
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 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 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']
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 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']
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']
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_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_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
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
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_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
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}
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_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(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}
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']
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']
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_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()
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_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'
async def delete(cls, access_key: str): """ Deletes an existing keypair with given ACCESSKEY. """ q = 'mutation($access_key: String!) {' \ ' delete_keypair(access_key: $access_key) {' \ ' ok msg' \ ' }' \ '}' variables = { 'access_key': access_key, } 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['delete_keypair']
async def delete(cls, gid: str): ''' Deletes an existing group. ''' query = textwrap.dedent('''\ mutation($gid: String!) { delete_group(gid: $gid) { ok msg } } ''') 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['delete_group']
async def list(cls, user_id: Union[int, str] = None, is_active: bool = None, fields: Iterable[str] = None) -> Sequence[dict]: ''' Lists the keypairs. You need an admin privilege for this operation. ''' if fields is None: fields = ( 'access_key', 'secret_key', 'is_active', 'is_admin', ) if user_id is None: q = 'query($is_active: Boolean) {' \ ' keypairs(is_active: $is_active) {' \ ' $fields' \ ' }' \ '}' else: uid_type = 'Int!' if isinstance(user_id, int) else 'String!' q = 'query($email: {0}, $is_active: Boolean) {{'.format(uid_type) + \ ' keypairs(email: $email, is_active: $is_active) {' \ ' $fields' \ ' }' \ '}' q = q.replace('$fields', ' '.join(fields)) variables = { 'is_active': is_active, } if user_id is not None: variables['email'] = user_id 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['keypairs']
async def create(cls, domain_name: str, name: str, description: str = '', is_active: bool = True, total_resource_slots: str = None, allowed_vfolder_hosts: Iterable[str] = None, integration_id: str = None, fields: Iterable[str] = None) -> dict: ''' Creates a new group with the given options. You need an admin privilege for this operation. ''' if fields is None: fields = ('id', 'domain_name', 'name',) query = textwrap.dedent('''\ mutation($name: String!, $input: GroupInput!) { create_group(name: $name, props: $input) { ok msg group {$fields} } } ''') query = query.replace('$fields', ' '.join(fields)) variables = { 'name': name, 'input': { 'description': description, 'is_active': is_active, 'domain_name': domain_name, '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['create_group']
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 create(cls, user_id: Union[int, str], is_active: bool = True, is_admin: bool = False, resource_policy: str = None, rate_limit: int = None, fields: Iterable[str] = None) -> dict: ''' Creates a new keypair with the given options. You need an admin privilege for this operation. ''' if fields is None: fields = ('access_key', 'secret_key') uid_type = 'Int!' if isinstance(user_id, int) else 'String!' q = 'mutation($user_id: {0}, $input: KeyPairInput!) {{'.format(uid_type) + \ ' create_keypair(user_id: $user_id, props: $input) {' \ ' ok msg keypair { $fields }' \ ' }' \ '}' q = q.replace('$fields', ' '.join(fields)) variables = { 'user_id': user_id, '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['create_keypair']
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(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_fetch_invalid_method_async(): async with AsyncSession() as session: rqst = Request('STRANGE', '/') with pytest.raises(AssertionError): async with rqst.fetch(): pass