Shore up multi tenant tests (#4484)

* update

* fix

* finalize`

* remove unnecessary prints

* fix

* k
This commit is contained in:
pablonyx
2025-04-14 11:34:57 -07:00
committed by GitHub
parent 68c6c1f4f8
commit e572ce95e7
10 changed files with 267 additions and 109 deletions

View File

@@ -78,15 +78,19 @@ class ChatSessionManager:
use_existing_user_message=use_existing_user_message,
)
headers = (
user_performing_action.headers
if user_performing_action
else GENERAL_HEADERS
)
cookies = user_performing_action.cookies if user_performing_action else None
response = requests.post(
f"{API_SERVER_URL}/chat/send-message",
json=chat_message_req.model_dump(),
headers=(
user_performing_action.headers
if user_performing_action
else GENERAL_HEADERS
),
headers=headers,
stream=True,
cookies=cookies,
)
return ChatSessionManager.analyze_response(response)

View File

@@ -6,7 +6,8 @@ INVITED_BASIC_USER = "basic_user"
INVITED_BASIC_USER_EMAIL = "basic_user@test.com"
def test_user_invitation_flow(reset_multitenant: None) -> None:
def test_admin_can_invite_users(reset_multitenant: None) -> None:
"""Test that an admin can invite both registered and non-registered users."""
# Create first user (admin)
admin_user: DATestUser = UserManager.create(name="admin")
assert UserManager.is_role(admin_user, UserRole.ADMIN)
@@ -19,16 +20,44 @@ def test_user_invitation_flow(reset_multitenant: None) -> None:
UserManager.invite_user(invited_user.email, admin_user)
UserManager.invite_user(INVITED_BASIC_USER_EMAIL, admin_user)
# Verify users are in the invited users list
invited_users = UserManager.get_invited_users(admin_user)
assert invited_user.email in [
user.email for user in invited_users
], f"User {invited_user.email} not found in invited users list"
def test_non_registered_user_gets_basic_role(reset_multitenant: None) -> None:
"""Test that a non-registered user gets a BASIC role when they register after being invited."""
# Create admin user
admin_user: DATestUser = UserManager.create(name="admin")
assert UserManager.is_role(admin_user, UserRole.ADMIN)
# Admin user invites a non-registered user
UserManager.invite_user(INVITED_BASIC_USER_EMAIL, admin_user)
# Non-registered user registers
invited_basic_user: DATestUser = UserManager.create(
name=INVITED_BASIC_USER, email=INVITED_BASIC_USER_EMAIL
)
assert UserManager.is_role(invited_basic_user, UserRole.BASIC)
# Verify the user is in the invited users list
invited_users = UserManager.get_invited_users(admin_user)
assert invited_user.email in [
user.email for user in invited_users
], f"User {invited_user.email} not found in invited users list"
def test_user_can_accept_invitation(reset_multitenant: None) -> None:
"""Test that a user can accept an invitation and join the organization with BASIC role."""
# Create admin user
admin_user: DATestUser = UserManager.create(name="admin")
assert UserManager.is_role(admin_user, UserRole.ADMIN)
# Create a user to be invited
invited_user_email = "invited_user@test.com"
# User registers with the same email as the invitation
invited_user: DATestUser = UserManager.create(
name="invited_user", email=invited_user_email
)
# Admin user invites the user
UserManager.invite_user(invited_user_email, admin_user)
# Get user info to check tenant information
user_info = UserManager.get_user_info(invited_user)
@@ -41,16 +70,17 @@ def test_user_invitation_flow(reset_multitenant: None) -> None:
)
assert invited_tenant_id is not None, "Expected to find an invitation tenant_id"
# User accepts invitation
UserManager.accept_invitation(invited_tenant_id, invited_user)
# Get updated user info after accepting invitation
updated_user_info = UserManager.get_user_info(invited_user)
# User needs to reauthenticate after accepting invitation
# Simulate this by creating a new user instance with the same credentials
authenticated_user: DATestUser = UserManager.create(
name="invited_user", email=invited_user_email
)
# Verify the user is no longer in the invited users list
updated_invited_users = UserManager.get_invited_users(admin_user)
assert invited_user.email not in [
user.email for user in updated_invited_users
], f"User {invited_user.email} should not be in invited users list after accepting"
# Get updated user info after accepting invitation and reauthenticating
updated_user_info = UserManager.get_user_info(authenticated_user)
# Verify the user has BASIC role in the organization
assert (
@@ -64,7 +94,7 @@ def test_user_invitation_flow(reset_multitenant: None) -> None:
# Check if the invited user is in the list of users with BASIC role
invited_user_emails = [user.email for user in user_page.items]
assert invited_user.email in invited_user_emails, (
f"User {invited_user.email} not found in the list of basic users "
assert invited_user_email in invited_user_emails, (
f"User {invited_user_email} not found in the list of basic users "
f"in the organization. Available users: {invited_user_emails}"
)

View File

@@ -1,3 +1,5 @@
from typing import Any
from onyx.db.models import UserRole
from tests.integration.common_utils.managers.api_key import APIKeyManager
from tests.integration.common_utils.managers.cc_pair import CCPairManager
@@ -11,12 +13,12 @@ from tests.integration.common_utils.test_models import DATestChatSession
from tests.integration.common_utils.test_models import DATestUser
def test_multi_tenant_access_control(reset_multitenant: None) -> None:
# Creating an admin user (first user created is automatically an admin and also proviions the tenant
def setup_test_tenants(reset_multitenant: None) -> dict[str, Any]:
"""Helper function to set up test tenants with documents and users."""
# Creating an admin user for Tenant 1
admin_user1: DATestUser = UserManager.create(
email="admin@onyx-test.com",
)
assert UserManager.is_role(admin_user1, UserRole.ADMIN)
# Create Tenant 2 and its Admin User
@@ -35,6 +37,16 @@ def test_multi_tenant_access_control(reset_multitenant: None) -> None:
api_key_1.headers.update(admin_user1.headers)
LLMProviderManager.create(user_performing_action=admin_user1)
# Create connectors for Tenant 2
cc_pair_2: DATestCCPair = CCPairManager.create_from_scratch(
user_performing_action=admin_user2,
)
api_key_2: DATestAPIKey = APIKeyManager.create(
user_performing_action=admin_user2,
)
api_key_2.headers.update(admin_user2.headers)
LLMProviderManager.create(user_performing_action=admin_user2)
# Seed documents for Tenant 1
cc_pair_1.documents = []
doc1_tenant1 = DocumentManager.seed_doc_with_content(
@@ -49,16 +61,6 @@ def test_multi_tenant_access_control(reset_multitenant: None) -> None:
)
cc_pair_1.documents.extend([doc1_tenant1, doc2_tenant1])
# Create connectors for Tenant 2
cc_pair_2: DATestCCPair = CCPairManager.create_from_scratch(
user_performing_action=admin_user2,
)
api_key_2: DATestAPIKey = APIKeyManager.create(
user_performing_action=admin_user2,
)
api_key_2.headers.update(admin_user2.headers)
LLMProviderManager.create(user_performing_action=admin_user2)
# Seed documents for Tenant 2
cc_pair_2.documents = []
doc1_tenant2 = DocumentManager.seed_doc_with_content(
@@ -84,21 +86,36 @@ def test_multi_tenant_access_control(reset_multitenant: None) -> None:
user_performing_action=admin_user2
)
return {
"admin_user1": admin_user1,
"admin_user2": admin_user2,
"chat_session1": chat_session1,
"chat_session2": chat_session2,
"tenant1_doc_ids": tenant1_doc_ids,
"tenant2_doc_ids": tenant2_doc_ids,
}
def test_tenant1_can_access_own_documents(reset_multitenant: None) -> None:
"""Test that Tenant 1 can access its own documents but not Tenant 2's."""
test_data = setup_test_tenants(reset_multitenant)
# User 1 sends a message and gets a response
response1 = ChatSessionManager.send_message(
chat_session_id=chat_session1.id,
chat_session_id=test_data["chat_session1"].id,
message="What is in Tenant 1's documents?",
user_performing_action=admin_user1,
user_performing_action=test_data["admin_user1"],
)
# Assert that the search tool was used
assert response1.tool_name == "run_search"
response_doc_ids = {doc["document_id"] for doc in response1.tool_result or []}
assert tenant1_doc_ids.issubset(
assert test_data["tenant1_doc_ids"].issubset(
response_doc_ids
), "Not all Tenant 1 document IDs are in the response"
assert not response_doc_ids.intersection(
tenant2_doc_ids
test_data["tenant2_doc_ids"]
), "Tenant 2 document IDs should not be in the response"
# Assert that the contents are correct
@@ -107,21 +124,28 @@ def test_multi_tenant_access_control(reset_multitenant: None) -> None:
for doc in response1.tool_result or []
), "Tenant 1 Document Content not found in any document"
def test_tenant2_can_access_own_documents(reset_multitenant: None) -> None:
"""Test that Tenant 2 can access its own documents but not Tenant 1's."""
test_data = setup_test_tenants(reset_multitenant)
# User 2 sends a message and gets a response
response2 = ChatSessionManager.send_message(
chat_session_id=chat_session2.id,
chat_session_id=test_data["chat_session2"].id,
message="What is in Tenant 2's documents?",
user_performing_action=admin_user2,
user_performing_action=test_data["admin_user2"],
)
# Assert that the search tool was used
assert response2.tool_name == "run_search"
# Assert that the tool_result contains Tenant 2's documents
response_doc_ids = {doc["document_id"] for doc in response2.tool_result or []}
assert tenant2_doc_ids.issubset(
assert test_data["tenant2_doc_ids"].issubset(
response_doc_ids
), "Not all Tenant 2 document IDs are in the response"
assert not response_doc_ids.intersection(
tenant1_doc_ids
test_data["tenant1_doc_ids"]
), "Tenant 1 document IDs should not be in the response"
# Assert that the contents are correct
@@ -130,28 +154,91 @@ def test_multi_tenant_access_control(reset_multitenant: None) -> None:
for doc in response2.tool_result or []
), "Tenant 2 Document Content not found in any document"
def test_tenant1_cannot_access_tenant2_documents(reset_multitenant: None) -> None:
"""Test that Tenant 1 cannot access Tenant 2's documents."""
test_data = setup_test_tenants(reset_multitenant)
# User 1 tries to access Tenant 2's documents
response_cross = ChatSessionManager.send_message(
chat_session_id=chat_session1.id,
chat_session_id=test_data["chat_session1"].id,
message="What is in Tenant 2's documents?",
user_performing_action=admin_user1,
user_performing_action=test_data["admin_user1"],
)
# Assert that the search tool was used
assert response_cross.tool_name == "run_search"
# Assert that the tool_result is empty or does not contain Tenant 2's documents
response_doc_ids = {doc["document_id"] for doc in response_cross.tool_result or []}
# Ensure none of Tenant 2's document IDs are in the response
assert not response_doc_ids.intersection(tenant2_doc_ids)
assert not response_doc_ids.intersection(test_data["tenant2_doc_ids"])
def test_tenant2_cannot_access_tenant1_documents(reset_multitenant: None) -> None:
"""Test that Tenant 2 cannot access Tenant 1's documents."""
test_data = setup_test_tenants(reset_multitenant)
# User 2 tries to access Tenant 1's documents
response_cross2 = ChatSessionManager.send_message(
chat_session_id=chat_session2.id,
chat_session_id=test_data["chat_session2"].id,
message="What is in Tenant 1's documents?",
user_performing_action=admin_user2,
user_performing_action=test_data["admin_user2"],
)
# Assert that the search tool was used
assert response_cross2.tool_name == "run_search"
# Assert that the tool_result is empty or does not contain Tenant 1's documents
response_doc_ids = {doc["document_id"] for doc in response_cross2.tool_result or []}
# Ensure none of Tenant 1's document IDs are in the response
assert not response_doc_ids.intersection(tenant1_doc_ids)
assert not response_doc_ids.intersection(test_data["tenant1_doc_ids"])
def test_multi_tenant_access_control(reset_multitenant: None) -> None:
"""Legacy test for multi-tenant access control."""
test_data = setup_test_tenants(reset_multitenant)
# User 1 sends a message and gets a response with only Tenant 1's documents
response1 = ChatSessionManager.send_message(
chat_session_id=test_data["chat_session1"].id,
message="What is in Tenant 1's documents?",
user_performing_action=test_data["admin_user1"],
)
assert response1.tool_name == "run_search"
response_doc_ids = {doc["document_id"] for doc in response1.tool_result or []}
assert test_data["tenant1_doc_ids"].issubset(response_doc_ids)
assert not response_doc_ids.intersection(test_data["tenant2_doc_ids"])
# User 2 sends a message and gets a response with only Tenant 2's documents
response2 = ChatSessionManager.send_message(
chat_session_id=test_data["chat_session2"].id,
message="What is in Tenant 2's documents?",
user_performing_action=test_data["admin_user2"],
)
assert response2.tool_name == "run_search"
response_doc_ids = {doc["document_id"] for doc in response2.tool_result or []}
assert test_data["tenant2_doc_ids"].issubset(response_doc_ids)
assert not response_doc_ids.intersection(test_data["tenant1_doc_ids"])
# User 1 tries to access Tenant 2's documents and fails
response_cross = ChatSessionManager.send_message(
chat_session_id=test_data["chat_session1"].id,
message="What is in Tenant 2's documents?",
user_performing_action=test_data["admin_user1"],
)
assert response_cross.tool_name == "run_search"
response_doc_ids = {doc["document_id"] for doc in response_cross.tool_result or []}
assert not response_doc_ids.intersection(test_data["tenant2_doc_ids"])
# User 2 tries to access Tenant 1's documents and fails
response_cross2 = ChatSessionManager.send_message(
chat_session_id=test_data["chat_session2"].id,
message="What is in Tenant 1's documents?",
user_performing_action=test_data["admin_user2"],
)
assert response_cross2.tool_name == "run_search"
response_doc_ids = {doc["document_id"] for doc in response_cross2.tool_result or []}
assert not response_doc_ids.intersection(test_data["tenant1_doc_ids"])

View File

@@ -8,12 +8,51 @@ from tests.integration.common_utils.managers.user import UserManager
from tests.integration.common_utils.test_models import DATestUser
# Test flow from creating tenant to registering as a user
def test_tenant_creation(reset_multitenant: None) -> None:
def test_first_user_is_admin(reset_multitenant: None) -> None:
"""Test that the first user of a tenant is automatically assigned ADMIN role."""
test_user: DATestUser = UserManager.create(name="test", email="test@test.com")
assert UserManager.is_role(test_user, UserRole.ADMIN)
def test_admin_can_create_credential(reset_multitenant: None) -> None:
"""Test that an admin user can create a credential in their tenant."""
# Create admin user
test_user: DATestUser = UserManager.create(name="test", email="test@test.com")
assert UserManager.is_role(test_user, UserRole.ADMIN)
# Create credential
test_credential = CredentialManager.create(
name="admin_test_credential",
source=DocumentSource.FILE,
curator_public=False,
user_performing_action=test_user,
)
assert test_credential is not None
def test_admin_can_create_connector(reset_multitenant: None) -> None:
"""Test that an admin user can create a connector in their tenant."""
# Create admin user
test_user: DATestUser = UserManager.create(name="test", email="test@test.com")
assert UserManager.is_role(test_user, UserRole.ADMIN)
# Create connector
test_connector = ConnectorManager.create(
name="admin_test_connector",
source=DocumentSource.FILE,
access_type=AccessType.PRIVATE,
user_performing_action=test_user,
)
assert test_connector is not None
def test_admin_can_create_and_verify_cc_pair(reset_multitenant: None) -> None:
"""Test that an admin user can create and verify a connector-credential pair in their tenant."""
# Create admin user
test_user: DATestUser = UserManager.create(name="test", email="test@test.com")
assert UserManager.is_role(test_user, UserRole.ADMIN)
# Create credential
test_credential = CredentialManager.create(
name="admin_test_credential",
source=DocumentSource.FILE,
@@ -21,6 +60,7 @@ def test_tenant_creation(reset_multitenant: None) -> None:
user_performing_action=test_user,
)
# Create connector
test_connector = ConnectorManager.create(
name="admin_test_connector",
source=DocumentSource.FILE,
@@ -28,6 +68,7 @@ def test_tenant_creation(reset_multitenant: None) -> None:
user_performing_action=test_user,
)
# Create cc_pair
test_cc_pair = CCPairManager.create(
connector_id=test_connector.id,
credential_id=test_credential.id,
@@ -35,5 +76,7 @@ def test_tenant_creation(reset_multitenant: None) -> None:
access_type=AccessType.PRIVATE,
user_performing_action=test_user,
)
assert test_cc_pair is not None
# Verify cc_pair
CCPairManager.verify(cc_pair=test_cc_pair, user_performing_action=test_user)

View File

@@ -29,12 +29,11 @@ function LLMProviderUpdateModal({
llmProviderDescriptor?.name ||
"Custom LLM Provider";
const hasAdvancedOptions = llmProviderDescriptor?.name != "azure";
return (
<Modal
title={`${llmProviderDescriptor ? "Configure" : "Setup"} ${providerName}`}
onOutsideClick={() => onClose()}
hideOverflow={true}
>
<div className="max-h-[70vh] overflow-y-auto px-4">
{llmProviderDescriptor ? (
@@ -44,7 +43,6 @@ function LLMProviderUpdateModal({
existingLlmProvider={existingLlmProvider}
shouldMarkAsDefault={shouldMarkAsDefault}
setPopup={setPopup}
hasAdvancedOptions={hasAdvancedOptions}
/>
) : (
<CustomLLMProviderUpdateForm

View File

@@ -35,10 +35,12 @@ function LLMProviderUpdateModal({
existingLlmProvider?.name ||
"Custom LLM Provider";
const hasAdvancedOptions = llmProviderDescriptor?.name != "azure";
return (
<Modal title={`Setup ${providerName}`} onOutsideClick={() => onClose()}>
<Modal
title={`Setup ${providerName}`}
onOutsideClick={() => onClose()}
hideOverflow={true}
>
<div className="max-h-[70vh] overflow-y-auto px-4">
{llmProviderDescriptor ? (
<LLMProviderUpdateForm
@@ -47,7 +49,6 @@ function LLMProviderUpdateModal({
existingLlmProvider={existingLlmProvider}
shouldMarkAsDefault={shouldMarkAsDefault}
setPopup={setPopup}
hasAdvancedOptions={hasAdvancedOptions}
/>
) : (
<CustomLLMProviderUpdateForm

View File

@@ -29,7 +29,6 @@ export function LLMProviderUpdateForm({
setPopup,
hideSuccess,
firstTimeConfiguration = false,
hasAdvancedOptions = false,
}: {
llmProviderDescriptor: WellKnownLLMProviderDescriptor;
onClose: () => void;
@@ -40,7 +39,6 @@ export function LLMProviderUpdateForm({
// Set this when this is the first time the user is setting Onyx up.
firstTimeConfiguration?: boolean;
hasAdvancedOptions?: boolean;
}) {
const { mutate } = useSWRConfig();
@@ -302,7 +300,7 @@ export function LLMProviderUpdateForm({
}
})}
{hasAdvancedOptions && !firstTimeConfiguration && (
{!firstTimeConfiguration && (
<>
<Separator />
@@ -364,52 +362,49 @@ export function LLMProviderUpdateForm({
/>
))}
{hasAdvancedOptions && (
<>
<Separator />
<AdvancedOptionsToggle
showAdvancedOptions={showAdvancedOptions}
setShowAdvancedOptions={setShowAdvancedOptions}
/>
{showAdvancedOptions && (
<>
{llmProviderDescriptor.llm_names.length > 0 && (
<div className="w-full">
<MultiSelectField
selectedInitially={
formikProps.values.display_model_names
}
name="display_model_names"
label="Display Models"
subtext="Select the models to make available to users. Unselected models will not be available."
options={llmProviderDescriptor.llm_names.map(
(name) => ({
value: name,
// don't clean up names here to give admins descriptive names / handle duplicates
// like us.anthropic.claude-3-7-sonnet-20250219-v1:0 and anthropic.claude-3-7-sonnet-20250219-v1:0
label: name,
})
)}
onChange={(selected) =>
formikProps.setFieldValue(
"display_model_names",
selected
)
}
/>
</div>
)}
<IsPublicGroupSelector
formikProps={formikProps}
objectName="LLM Provider"
publicToWhom="all users"
enforceGroupSelection={true}
/>
</>
)}
</>
)}
<>
<Separator />
<AdvancedOptionsToggle
showAdvancedOptions={showAdvancedOptions}
setShowAdvancedOptions={setShowAdvancedOptions}
/>
{showAdvancedOptions && (
<>
{llmProviderDescriptor.llm_names.length > 0 && (
<div className="w-full">
<MultiSelectField
selectedInitially={
formikProps.values.display_model_names
}
name="display_model_names"
label="Display Models"
subtext="Select the models to make available to users. Unselected models will not be available."
options={llmProviderDescriptor.llm_names.map(
(name) => ({
value: name,
// don't clean up names here to give admins descriptive names / handle duplicates
// like us.anthropic.claude-3-7-sonnet-20250219-v1:0 and anthropic.claude-3-7-sonnet-20250219-v1:0
label: name,
})
)}
onChange={(selected) =>
formikProps.setFieldValue(
"display_model_names",
selected
)
}
/>
</div>
)}
<IsPublicGroupSelector
formikProps={formikProps}
objectName="LLM Provider"
publicToWhom="Users"
enforceGroupSelection={true}
/>
</>
)}
</>
</>
)}

View File

@@ -74,7 +74,6 @@ export const IsPublicGroupSelector = <T extends IsPublicGroupSelectorFormType>({
return (
<div>
<Separator />
{isAdmin && (
<>
<BooleanFormField

View File

@@ -24,6 +24,7 @@ interface ModalProps {
removeBottomPadding?: boolean;
removePadding?: boolean;
increasedPadding?: boolean;
hideOverflow?: boolean;
}
export function Modal({
@@ -43,6 +44,7 @@ export function Modal({
removeBottomPadding,
removePadding,
increasedPadding,
hideOverflow,
}: ModalProps) {
const modalRef = useRef<HTMLDivElement>(null);
const [isMounted, setIsMounted] = useState(false);
@@ -92,7 +94,7 @@ export function Modal({
flex
flex-col
${heightOverride ? `h-${heightOverride}` : "max-h-[90vh]"}
overflow-auto
${hideOverflow ? "overflow-hidden" : "overflow-auto"}
`}
>
{onOutsideClick && !hideCloseButton && (

View File

@@ -174,7 +174,6 @@ export const CustomTooltip = ({
: {}
}
>
lll
{content}
</div>
</div>,