mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-05-13 13:20:15 +02:00
87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
from datetime import datetime
|
|
from datetime import timedelta
|
|
|
|
import jwt
|
|
import requests
|
|
|
|
from onyx.server.manage.models import AllUsersResponse
|
|
from onyx.server.models import FullUserSnapshot
|
|
from onyx.server.models import InvitedUserSnapshot
|
|
from tests.integration.common_utils.constants import API_SERVER_URL
|
|
from tests.integration.common_utils.constants import GENERAL_HEADERS
|
|
from tests.integration.common_utils.test_models import DATestUser
|
|
|
|
|
|
def generate_auth_token() -> str:
|
|
payload = {
|
|
"iss": "control_plane",
|
|
"exp": datetime.utcnow() + timedelta(minutes=5),
|
|
"iat": datetime.utcnow(),
|
|
"scope": "tenant:create",
|
|
}
|
|
token = jwt.encode(payload, "", algorithm="HS256")
|
|
return token
|
|
|
|
|
|
class TenantManager:
|
|
@staticmethod
|
|
def create(
|
|
tenant_id: str | None = None,
|
|
initial_admin_email: str | None = None,
|
|
referral_source: str | None = None,
|
|
) -> dict[str, str]:
|
|
body = {
|
|
"tenant_id": tenant_id,
|
|
"initial_admin_email": initial_admin_email,
|
|
"referral_source": referral_source,
|
|
}
|
|
|
|
token = generate_auth_token()
|
|
headers = {
|
|
"Authorization": f"Bearer {token}",
|
|
"X-API-KEY": "",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
response = requests.post(
|
|
url=f"{API_SERVER_URL}/tenants/create",
|
|
json=body,
|
|
headers=headers,
|
|
)
|
|
|
|
response.raise_for_status()
|
|
|
|
return response.json()
|
|
|
|
@staticmethod
|
|
def get_all_users(
|
|
user_performing_action: DATestUser | None = None,
|
|
) -> AllUsersResponse:
|
|
response = requests.get(
|
|
url=f"{API_SERVER_URL}/manage/users",
|
|
headers=user_performing_action.headers
|
|
if user_performing_action
|
|
else GENERAL_HEADERS,
|
|
)
|
|
response.raise_for_status()
|
|
|
|
data = response.json()
|
|
return AllUsersResponse(
|
|
accepted=[FullUserSnapshot(**user) for user in data["accepted"]],
|
|
invited=[InvitedUserSnapshot(**user) for user in data["invited"]],
|
|
slack_users=[FullUserSnapshot(**user) for user in data["slack_users"]],
|
|
accepted_pages=data["accepted_pages"],
|
|
invited_pages=data["invited_pages"],
|
|
slack_users_pages=data["slack_users_pages"],
|
|
)
|
|
|
|
@staticmethod
|
|
def verify_user_in_tenant(
|
|
user: DATestUser, user_performing_action: DATestUser | None = None
|
|
) -> None:
|
|
all_users = TenantManager.get_all_users(user_performing_action)
|
|
for accepted_user in all_users.accepted:
|
|
if accepted_user.email == user.email and accepted_user.id == user.id:
|
|
return
|
|
raise ValueError(f"User {user.email} not found in tenant")
|