from copy import deepcopy from urllib.parse import urlencode from uuid import uuid4 import requests from danswer.db.models import UserRole from danswer.server.manage.models import AllUsersResponse from danswer.server.models import FullUserSnapshot from danswer.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 class UserManager: @staticmethod def create( name: str | None = None, email: str | None = None, ) -> DATestUser: if name is None: name = f"test{str(uuid4())}" if email is None: email = f"{name}@test.com" password = "test" body = { "email": email, "username": email, "password": password, } response = requests.post( url=f"{API_SERVER_URL}/auth/register", json=body, headers=GENERAL_HEADERS, ) response.raise_for_status() test_user = DATestUser( id=response.json()["id"], email=email, password=password, headers=deepcopy(GENERAL_HEADERS), ) print(f"Created user {test_user.email}") return UserManager.login_as_user(test_user) @staticmethod def login_as_user(test_user: DATestUser) -> DATestUser: data = urlencode( { "username": test_user.email, "password": test_user.password, } ) headers = test_user.headers.copy() headers["Content-Type"] = "application/x-www-form-urlencoded" response = requests.post( url=f"{API_SERVER_URL}/auth/login", data=data, headers=headers, ) response.raise_for_status() result_cookie = next(iter(response.cookies), None) if not result_cookie: raise Exception("Failed to login") print(f"Logged in as {test_user.email}") cookie = f"{result_cookie.name}={result_cookie.value}" test_user.headers["Cookie"] = cookie return test_user @staticmethod def verify_role(user_to_verify: DATestUser, target_role: UserRole) -> bool: response = requests.get( url=f"{API_SERVER_URL}/me", headers=user_to_verify.headers, ) response.raise_for_status() return target_role == UserRole(response.json().get("role", "")) @staticmethod def set_role( user_to_set: DATestUser, target_role: UserRole, user_to_perform_action: DATestUser | None = None, ) -> None: if user_to_perform_action is None: user_to_perform_action = user_to_set response = requests.patch( url=f"{API_SERVER_URL}/manage/set-user-role", json={"user_email": user_to_set.email, "new_role": target_role.value}, headers=user_to_perform_action.headers, ) response.raise_for_status() @staticmethod def verify( user: DATestUser, user_to_perform_action: DATestUser | None = None ) -> None: if user_to_perform_action is None: user_to_perform_action = user response = requests.get( url=f"{API_SERVER_URL}/manage/users", headers=user_to_perform_action.headers if user_to_perform_action else GENERAL_HEADERS, ) response.raise_for_status() data = response.json() all_users = AllUsersResponse( accepted=[FullUserSnapshot(**user) for user in data["accepted"]], invited=[InvitedUserSnapshot(**user) for user in data["invited"]], accepted_pages=data["accepted_pages"], invited_pages=data["invited_pages"], ) 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")