Skip to content

Commit

Permalink
Global Curator Fix + Testing (#2591)
Browse files Browse the repository at this point in the history
* Global Curator Fix

* test fix
  • Loading branch information
hagen-danswer authored Sep 28, 2024
1 parent 50dd3c8 commit 1cff2b8
Show file tree
Hide file tree
Showing 3 changed files with 91 additions and 1 deletion.
6 changes: 5 additions & 1 deletion backend/ee/danswer/db/user_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,8 +151,12 @@ def validate_user_creation_permissions(
status_code=400,
detail=detail,
)

user_curated_groups = fetch_user_groups_for_user(
db_session=db_session, user_id=user.id, only_curator_groups=True
db_session=db_session,
user_id=user.id,
# Global curators can curate all groups they are member of
only_curator_groups=user.role != UserRole.GLOBAL_CURATOR,
)
user_curated_group_ids = set([group.id for group in user_curated_groups])
target_group_ids_set = set(target_group_ids)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -85,3 +85,84 @@ def test_whole_curator_flow(reset: None) -> None:
verify_deleted=True,
user_performing_action=admin_user,
)


def test_global_curator_flow(reset: None) -> None:
# Creating an admin user (first user created is automatically an admin)
admin_user: DATestUser = UserManager.create(name="admin_user")
assert UserManager.verify_role(admin_user, UserRole.ADMIN)

# Creating a user
global_curator: DATestUser = UserManager.create(name="global_curator")
assert UserManager.verify_role(global_curator, UserRole.BASIC)

# Set the user to a global curator
UserManager.set_role(
user_to_set=global_curator,
target_role=UserRole.GLOBAL_CURATOR,
user_to_perform_action=admin_user,
)
assert UserManager.verify_role(global_curator, UserRole.GLOBAL_CURATOR)

# Creating a user group containing the global curator
user_group_1 = UserGroupManager.create(
name="user_group_1",
user_ids=[global_curator.id],
cc_pair_ids=[],
user_performing_action=admin_user,
)
UserGroupManager.wait_for_sync(
user_groups_to_check=[user_group_1], user_performing_action=admin_user
)

# Creating a credential as global curator
test_credential = CredentialManager.create(
name="curator_test_credential",
source=DocumentSource.FILE,
curator_public=False,
groups=[user_group_1.id],
user_performing_action=global_curator,
)

# Creating a connector as global curator
test_connector = ConnectorManager.create(
name="curator_test_connector",
source=DocumentSource.FILE,
is_public=False,
groups=[user_group_1.id],
user_performing_action=global_curator,
)

# Test editing the connector
test_connector.name = "updated_test_connector"
ConnectorManager.edit(
connector=test_connector, user_performing_action=global_curator
)

# Creating a CC pair as global curator
test_cc_pair = CCPairManager.create(
connector_id=test_connector.id,
credential_id=test_credential.id,
name="curator_test_cc_pair",
access_type=AccessType.PRIVATE,
groups=[user_group_1.id],
user_performing_action=global_curator,
)

CCPairManager.verify(cc_pair=test_cc_pair, user_performing_action=admin_user)

# Verify that the curator can pause and unpause the CC pair
CCPairManager.pause_cc_pair(
cc_pair=test_cc_pair, user_performing_action=global_curator
)

# Verify that the curator can delete the CC pair
CCPairManager.delete(cc_pair=test_cc_pair, user_performing_action=global_curator)
CCPairManager.wait_for_deletion_completion(user_performing_action=global_curator)

# Verify that the CC pair has been deleted
CCPairManager.verify(
cc_pair=test_cc_pair,
verify_deleted=True,
user_performing_action=admin_user,
)
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ def test_user_group_deletion(reset: None, vespa_client: vespa_fixture) -> None:
user_performing_action=admin_user,
)

# Create other objects that are related to the user group
credential: DATestCredential = CredentialManager.create(
groups=[user_group.id],
user_performing_action=admin_user,
Expand All @@ -81,6 +82,7 @@ def test_user_group_deletion(reset: None, vespa_client: vespa_fixture) -> None:
user_performing_action=admin_user,
)

# Delete the user group
UserGroupManager.delete(
user_group=user_group,
user_performing_action=admin_user,
Expand All @@ -90,10 +92,13 @@ def test_user_group_deletion(reset: None, vespa_client: vespa_fixture) -> None:
user_groups_to_check=[user_group], user_performing_action=admin_user
)

# Set our expected local representations to empty
credential.groups = []
document_set.groups = []
llm_provider.groups = []
persona.groups = []

# Verify that the local representations were updated
CredentialManager.verify(
credential=credential,
user_performing_action=admin_user,
Expand Down

0 comments on commit 1cff2b8

Please sign in to comment.