259 lines
7.8 KiB
Python
259 lines
7.8 KiB
Python
import os
|
|
from dataclasses import dataclass
|
|
from unittest import TestCase
|
|
|
|
os.environ.setdefault('AWS_PROFILE', 'playground')
|
|
os.environ.setdefault('MULTI_USER_TABLE', 'MultiUserTable')
|
|
os_key_id = 'alias/sd-extension-password-key'
|
|
os.environ.setdefault('KEY_ID', os_key_id)
|
|
|
|
|
|
@dataclass
|
|
class MockContext:
|
|
aws_request_id: str
|
|
|
|
|
|
class InferenceApiTest(TestCase):
|
|
|
|
def test_kms(self):
|
|
from common.ddb_service.client import DynamoDbUtilsService
|
|
from libs.utils import KeyEncryptService
|
|
|
|
text = 'this text need to be encrypted'
|
|
key_id = os_key_id
|
|
|
|
key_client = KeyEncryptService()
|
|
cipher_text = key_client.encrypt(key_id, text)
|
|
ddb_service = DynamoDbUtilsService()
|
|
table_name = 'MultiUserTable'
|
|
from libs.data_types import User
|
|
ddb_service.put_items(table_name, User(
|
|
kind='user',
|
|
sort_key='superman',
|
|
password=cipher_text,
|
|
creator='alvindaiyan',
|
|
roles=['IT Operator', 'Designer']
|
|
).__dict__)
|
|
|
|
row = ddb_service.get_item(table_name, {
|
|
'kind': 'user',
|
|
'sort_key': 'alvindaiyan'
|
|
})
|
|
|
|
plain_txt = key_client.decrypt(key_id=key_id, cipher_text=row['password'].value)
|
|
assert plain_txt.decode() == text
|
|
|
|
def test_add_two_roles(self):
|
|
from common.roles_api import upsert_role
|
|
rolenames = ['IT Operator', 'Designer']
|
|
for rn in rolenames:
|
|
event = {
|
|
'role_name': rn,
|
|
'permissions': [
|
|
'train:all',
|
|
'checkpoint:all',
|
|
'inference:all',
|
|
'sagemaker_endpoint:all',
|
|
'user:all'
|
|
],
|
|
'creator': 'alvindaiyan'
|
|
}
|
|
resp = upsert_role(event, {})
|
|
print(resp)
|
|
|
|
def test_add_role(self):
|
|
from multi_users.roles_api import upsert_role
|
|
event = {
|
|
'role_name': f'avengers team',
|
|
'permissions': [
|
|
'train:all',
|
|
'checkpoint:all',
|
|
'inference:all',
|
|
'sagemaker_endpoint:all',
|
|
'user:all'
|
|
],
|
|
'creator': 'nickfury'
|
|
}
|
|
resp = upsert_role(event, {})
|
|
print(resp)
|
|
|
|
def test_add_roles(self, count):
|
|
from multi_users.roles_api import upsert_role
|
|
event = {
|
|
'role_name': f'RandomRole{count}',
|
|
'permissions': [
|
|
'train:all',
|
|
'checkpoint:all',
|
|
'inference:all',
|
|
'sagemaker_endpoint:all',
|
|
'user:all'
|
|
],
|
|
'creator': 'nickfury'
|
|
}
|
|
resp = upsert_role(event, {})
|
|
print(resp)
|
|
|
|
def test_batch_add_roles(self):
|
|
for i in range(100):
|
|
self.test_add_roles(i)
|
|
|
|
def test_add_user(self, count):
|
|
from multi_users.multi_users_api import upsert_user
|
|
event = {'body': {
|
|
'username': f'batman{count}',
|
|
'password': 'password',
|
|
'roles': ['IT Operator'],
|
|
'creator': 'spiderman'
|
|
}}
|
|
resp = upsert_user(event, {})
|
|
print(resp)
|
|
assert resp['status'] == 200
|
|
|
|
def test_add_user_with_no_permission(self):
|
|
self.test_add_user(0)
|
|
|
|
def test_add_user(self):
|
|
from multi_users.multi_users_api import upsert_user
|
|
event = {'body': {
|
|
'username': f'nickfury',
|
|
'password': 'password',
|
|
'roles': ['avengers manager'],
|
|
'creator': 'admin'
|
|
}}
|
|
resp = upsert_user(event, {})
|
|
print(resp)
|
|
assert resp['status'] == 200
|
|
|
|
def test_batch_add_users(self):
|
|
for i in range(300):
|
|
self.test_add_user(i)
|
|
|
|
def test_list_users_token(self):
|
|
from multi_users.multi_users_api import list_user
|
|
event = {
|
|
'queryStringParameters': {
|
|
# 'username': 'alvindaiyan',
|
|
'limit': 3,
|
|
'show_password': False,
|
|
},
|
|
'x-auth': {
|
|
'username': 'admin',
|
|
'role': ''
|
|
}
|
|
}
|
|
resp = list_user(event, {})
|
|
print(len(resp['users']))
|
|
print([u['username'] for u in resp['users']])
|
|
print(resp['last_evaluated_key'])
|
|
while resp['last_evaluated_key']:
|
|
event['queryStringParameters']['last_evaluated_key'] = resp['last_evaluated_key']
|
|
resp = list_user(event, {})
|
|
print(len(resp['users']))
|
|
print([u['username'] for u in resp['users']])
|
|
|
|
def test_list_all_users(self):
|
|
from multi_users.multi_users_api import list_user
|
|
event = {
|
|
'queryStringParameters': {
|
|
'show_password': False,
|
|
},
|
|
'x-auth': {
|
|
'username': 'thor',
|
|
'role': ''
|
|
}
|
|
}
|
|
resp = list_user(event, {})
|
|
print(len(resp['users']))
|
|
print([u['username'] for u in resp['users']])
|
|
|
|
def test_list_username(self):
|
|
from multi_users.multi_users_api import list_user
|
|
username = 'admin'
|
|
event = {
|
|
'queryStringParameters': {
|
|
'show_password': True,
|
|
'username': username
|
|
}
|
|
}
|
|
resp = list_user(event, {})
|
|
assert len(resp['users']) == 1
|
|
assert [u['username'] for u in resp['users']][0] == username
|
|
print(resp)
|
|
|
|
def test_list_role_token(self):
|
|
from multi_users.roles_api import list_roles
|
|
event = {
|
|
'queryStringParameters': {
|
|
'limit': 10,
|
|
}
|
|
}
|
|
resp = list_roles(event, {})
|
|
print(len(resp['roles']))
|
|
print([u['role_name'] for u in resp['roles']])
|
|
print(resp['last_evaluated_key'])
|
|
while resp['last_evaluated_key']:
|
|
event['queryStringParameters']['last_evaluated_key'] = resp['last_evaluated_key']
|
|
resp = list_roles(event, {})
|
|
print(len(resp['roles']))
|
|
print([u['role_name'] for u in resp['roles']])
|
|
|
|
def test_list_all_roles(self):
|
|
from multi_users.roles_api import list_roles
|
|
event = {
|
|
'queryStringParameters': {
|
|
},
|
|
'x-auth': {
|
|
'username': 'xman',
|
|
'role': ''
|
|
}
|
|
}
|
|
resp = list_roles(event, {})
|
|
print(len(resp['roles']))
|
|
print([u['role_name'] for u in resp['roles']])
|
|
|
|
def test_list_role(self):
|
|
from multi_users.roles_api import list_roles
|
|
role = 'RandomRole99'
|
|
event = {
|
|
'queryStringParameters': {
|
|
'role': role
|
|
},
|
|
|
|
}
|
|
resp = list_roles(event, {})
|
|
assert len(resp['roles']) == 1
|
|
assert [r['role_name'] for r in resp['roles']][0] == role
|
|
print(resp)
|
|
|
|
def test_delete_user(self):
|
|
from multi_users.multi_users_api import delete_user
|
|
username = 'batman10'
|
|
event = {
|
|
'pathStringParameters': {
|
|
'username': username
|
|
}
|
|
}
|
|
resp = delete_user(event, {})
|
|
print(resp)
|
|
|
|
from multi_users.multi_users_api import list_user
|
|
|
|
event = {
|
|
'queryStringParameters': {
|
|
'show_password': True,
|
|
'username': username
|
|
}
|
|
}
|
|
resp = list_user(event, {})
|
|
assert len(resp['users']) == 0
|
|
|
|
def test_get_user_roles(self):
|
|
from multi_users.utils import get_user_roles
|
|
import logging
|
|
logger = logging.getLogger('roles_api')
|
|
from common.ddb_service.client import DynamoDbUtilsService
|
|
|
|
ddb_service = DynamoDbUtilsService(logger=logger)
|
|
resp = get_user_roles(ddb_service, user_table_name='MultiUserTable', username='mickey')
|
|
print(resp)
|