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
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
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()