From 1cff2b82fdafad3bcb877d32b6baabbe29976af3 Mon Sep 17 00:00:00 2001 From: hagen-danswer Date: Sat, 28 Sep 2024 13:14:39 -0700 Subject: [PATCH] Global Curator Fix + Testing (#2591) * Global Curator Fix * test fix --- backend/ee/danswer/db/user_group.py | 6 +- .../permissions/test_whole_curator_flow.py | 81 +++++++++++++++++++ .../usergroup/test_user_group_deletion.py | 5 ++ 3 files changed, 91 insertions(+), 1 deletion(-) diff --git a/backend/ee/danswer/db/user_group.py b/backend/ee/danswer/db/user_group.py index 00ad159c1..863b9170e 100644 --- a/backend/ee/danswer/db/user_group.py +++ b/backend/ee/danswer/db/user_group.py @@ -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) diff --git a/backend/tests/integration/tests/permissions/test_whole_curator_flow.py b/backend/tests/integration/tests/permissions/test_whole_curator_flow.py index 2fbfe7b1c..1ce9052c1 100644 --- a/backend/tests/integration/tests/permissions/test_whole_curator_flow.py +++ b/backend/tests/integration/tests/permissions/test_whole_curator_flow.py @@ -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, + ) diff --git a/backend/tests/integration/tests/usergroup/test_user_group_deletion.py b/backend/tests/integration/tests/usergroup/test_user_group_deletion.py index eed0d9d8c..31d7f3d20 100644 --- a/backend/tests/integration/tests/usergroup/test_user_group_deletion.py +++ b/backend/tests/integration/tests/usergroup/test_user_group_deletion.py @@ -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, @@ -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, @@ -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,