diff --git a/backend/tests/integration/common_utils/managers/document_search.py b/backend/tests/integration/common_utils/managers/document_search.py new file mode 100644 index 000000000..4fe2442b6 --- /dev/null +++ b/backend/tests/integration/common_utils/managers/document_search.py @@ -0,0 +1,39 @@ +import requests + +from danswer.search.enums import LLMEvaluationType +from danswer.search.enums import SearchType +from danswer.search.models import RetrievalDetails +from danswer.search.models import SavedSearchDocWithContent +from ee.danswer.server.query_and_chat.models import DocumentSearchRequest +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 DocumentSearchManager: + @staticmethod + def search_documents( + query: str, + search_type: SearchType = SearchType.KEYWORD, + user_performing_action: DATestUser | None = None, + ) -> list[str]: + search_request = DocumentSearchRequest( + message=query, + search_type=search_type, + retrieval_options=RetrievalDetails(), + evaluation_type=LLMEvaluationType.SKIP, + ) + result = requests.post( + url=f"{API_SERVER_URL}/query/document-search", + json=search_request.model_dump(), + headers=user_performing_action.headers + if user_performing_action + else GENERAL_HEADERS, + ) + result.raise_for_status() + result_json = result.json() + top_documents: list[SavedSearchDocWithContent] = [ + SavedSearchDocWithContent(**doc) for doc in result_json["top_documents"] + ] + document_content_list: list[str] = [doc.content for doc in top_documents] + return document_content_list diff --git a/backend/tests/integration/connector_job_tests/slack/test_permission_sync.py b/backend/tests/integration/connector_job_tests/slack/test_permission_sync.py index 9114c2dde..7fdb1a042 100644 --- a/backend/tests/integration/connector_job_tests/slack/test_permission_sync.py +++ b/backend/tests/integration/connector_job_tests/slack/test_permission_sync.py @@ -4,19 +4,16 @@ from datetime import timezone from typing import Any import pytest -import requests from danswer.connectors.models import InputType from danswer.db.enums import AccessType -from danswer.search.enums import LLMEvaluationType -from danswer.search.enums import SearchType -from danswer.search.models import RetrievalDetails from danswer.server.documents.models import DocumentSource -from ee.danswer.server.query_and_chat.models import DocumentSearchRequest -from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.connector import ConnectorManager from tests.integration.common_utils.managers.credential import CredentialManager +from tests.integration.common_utils.managers.document_search import ( + DocumentSearchManager, +) from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestCCPair @@ -27,7 +24,7 @@ from tests.integration.common_utils.vespa import vespa_fixture from tests.integration.connector_job_tests.slack.slack_api_utils import SlackManager -@pytest.mark.skip(reason="flaky - see DAN-789 for example") +@pytest.mark.xfail(reason="flaky - see DAN-789 for example", strict=False) def test_slack_permission_sync( reset: None, vespa_client: vespa_fixture, @@ -100,6 +97,7 @@ def test_slack_permission_sync( public_message = "Steve's favorite number is 809752" private_message = "Sara's favorite number is 346794" + # Add messages to channels SlackManager.add_message_to_channel( slack_client=slack_client, channel=public_channel, @@ -133,42 +131,20 @@ def test_slack_permission_sync( ) # Search as admin with access to both channels - search_request = DocumentSearchRequest( - message="favorite number", - search_type=SearchType.KEYWORD, - retrieval_options=RetrievalDetails(), - evaluation_type=LLMEvaluationType.SKIP, + danswer_doc_message_strings = DocumentSearchManager.search_documents( + query="favorite number", + user_performing_action=admin_user, ) - search_request_body = search_request.model_dump() - result = requests.post( - url=f"{API_SERVER_URL}/query/document-search", - json=search_request_body, - headers=admin_user.headers, - ) - result.raise_for_status() - found_docs = result.json()["top_documents"] - danswer_doc_message_strings = [doc["content"] for doc in found_docs] # Ensure admin user can see messages from both channels assert public_message in danswer_doc_message_strings assert private_message in danswer_doc_message_strings # Search as test_user_2 with access to only the public channel - search_request = DocumentSearchRequest( - message="favorite number", - search_type=SearchType.KEYWORD, - retrieval_options=RetrievalDetails(), - evaluation_type=LLMEvaluationType.SKIP, + danswer_doc_message_strings = DocumentSearchManager.search_documents( + query="favorite number", + user_performing_action=test_user_2, ) - search_request_body = search_request.model_dump() - result = requests.post( - url=f"{API_SERVER_URL}/query/document-search", - json=search_request_body, - headers=test_user_2.headers, - ) - result.raise_for_status() - found_docs = result.json()["top_documents"] - danswer_doc_message_strings = [doc["content"] for doc in found_docs] print( "\ntop_documents content before removing from private channel for test_user_2: ", danswer_doc_message_strings, @@ -179,21 +155,10 @@ def test_slack_permission_sync( assert private_message not in danswer_doc_message_strings # Search as test_user_1 with access to both channels - search_request = DocumentSearchRequest( - message="favorite number", - search_type=SearchType.KEYWORD, - retrieval_options=RetrievalDetails(), - evaluation_type=LLMEvaluationType.SKIP, + danswer_doc_message_strings = DocumentSearchManager.search_documents( + query="favorite number", + user_performing_action=test_user_1, ) - search_request_body = search_request.model_dump() - result = requests.post( - url=f"{API_SERVER_URL}/query/document-search", - json=search_request_body, - headers=test_user_1.headers, - ) - result.raise_for_status() - found_docs = result.json()["top_documents"] - danswer_doc_message_strings = [doc["content"] for doc in found_docs] print( "\ntop_documents content before removing from private channel for test_user_1: ", danswer_doc_message_strings, @@ -228,21 +193,10 @@ def test_slack_permission_sync( # ----------------------------VERIFY THE CHANGES--------------------------- # Ensure test_user_1 can no longer see messages from the private channel # Search as test_user_1 with access to only the public channel - search_request = DocumentSearchRequest( - message="favorite number", - search_type=SearchType.KEYWORD, - retrieval_options=RetrievalDetails(), - evaluation_type=LLMEvaluationType.SKIP, + danswer_doc_message_strings = DocumentSearchManager.search_documents( + query="favorite number", + user_performing_action=test_user_1, ) - search_request_body = search_request.model_dump() - result = requests.post( - url=f"{API_SERVER_URL}/query/document-search", - json=search_request_body, - headers=test_user_1.headers, - ) - result.raise_for_status() - found_docs = result.json()["top_documents"] - danswer_doc_message_strings = [doc["content"] for doc in found_docs] print( "\ntop_documents content after removing from private channel for test_user_1: ", danswer_doc_message_strings, diff --git a/backend/tests/integration/connector_job_tests/slack/test_prune.py b/backend/tests/integration/connector_job_tests/slack/test_prune.py index 12cf7da7a..3410a477c 100644 --- a/backend/tests/integration/connector_job_tests/slack/test_prune.py +++ b/backend/tests/integration/connector_job_tests/slack/test_prune.py @@ -3,19 +3,15 @@ from datetime import datetime from datetime import timezone from typing import Any -import requests - from danswer.connectors.models import InputType from danswer.db.enums import AccessType -from danswer.search.enums import LLMEvaluationType -from danswer.search.enums import SearchType -from danswer.search.models import RetrievalDetails from danswer.server.documents.models import DocumentSource -from ee.danswer.server.query_and_chat.models import DocumentSearchRequest -from tests.integration.common_utils.constants import API_SERVER_URL from tests.integration.common_utils.managers.cc_pair import CCPairManager from tests.integration.common_utils.managers.connector import ConnectorManager from tests.integration.common_utils.managers.credential import CredentialManager +from tests.integration.common_utils.managers.document_search import ( + DocumentSearchManager, +) from tests.integration.common_utils.managers.llm_provider import LLMProviderManager from tests.integration.common_utils.managers.user import UserManager from tests.integration.common_utils.test_models import DATestCCPair @@ -134,21 +130,10 @@ def test_slack_prune( # ----------------------TEST THE SETUP-------------------------- # Search as admin with access to both channels - search_request = DocumentSearchRequest( - message="favorite number", - search_type=SearchType.KEYWORD, - retrieval_options=RetrievalDetails(), - evaluation_type=LLMEvaluationType.SKIP, + danswer_doc_message_strings = DocumentSearchManager.search_documents( + query="favorite number", + user_performing_action=admin_user, ) - search_request_body = search_request.model_dump() - result = requests.post( - url=f"{API_SERVER_URL}/query/document-search", - json=search_request_body, - headers=admin_user.headers, - ) - result.raise_for_status() - found_docs = result.json()["top_documents"] - danswer_doc_message_strings = [doc["content"] for doc in found_docs] print( "\ntop_documents content before deleting for admin: ", danswer_doc_message_strings, @@ -160,21 +145,10 @@ def test_slack_prune( assert message_to_delete in danswer_doc_message_strings # Search as test_user_1 with access to both channels - search_request = DocumentSearchRequest( - message="favorite number", - search_type=SearchType.KEYWORD, - retrieval_options=RetrievalDetails(), - evaluation_type=LLMEvaluationType.SKIP, + danswer_doc_message_strings = DocumentSearchManager.search_documents( + query="favorite number", + user_performing_action=test_user_1, ) - search_request_body = search_request.model_dump() - result = requests.post( - url=f"{API_SERVER_URL}/query/document-search", - json=search_request_body, - headers=test_user_1.headers, - ) - result.raise_for_status() - found_docs = result.json()["top_documents"] - danswer_doc_message_strings = [doc["content"] for doc in found_docs] print( "\ntop_documents content before deleting for test_user_1: ", danswer_doc_message_strings, @@ -202,21 +176,10 @@ def test_slack_prune( # ----------------------------VERIFY THE CHANGES--------------------------- # Ensure admin user can't see deleted messages # Search as admin user with access to only the public channel - search_request = DocumentSearchRequest( - message="favorite number", - search_type=SearchType.KEYWORD, - retrieval_options=RetrievalDetails(), - evaluation_type=LLMEvaluationType.SKIP, + danswer_doc_message_strings = DocumentSearchManager.search_documents( + query="favorite number", + user_performing_action=admin_user, ) - search_request_body = search_request.model_dump() - result = requests.post( - url=f"{API_SERVER_URL}/query/document-search", - json=search_request_body, - headers=admin_user.headers, - ) - result.raise_for_status() - found_docs = result.json()["top_documents"] - danswer_doc_message_strings = [doc["content"] for doc in found_docs] print( "\ntop_documents content after deleting for admin: ", danswer_doc_message_strings, @@ -229,21 +192,10 @@ def test_slack_prune( # Ensure test_user_1 can't see deleted messages # Search as test_user_1 with access to only the public channel - search_request = DocumentSearchRequest( - message="favorite number", - search_type=SearchType.KEYWORD, - retrieval_options=RetrievalDetails(), - evaluation_type=LLMEvaluationType.SKIP, + danswer_doc_message_strings = DocumentSearchManager.search_documents( + query="favorite number", + user_performing_action=test_user_1, ) - search_request_body = search_request.model_dump() - result = requests.post( - url=f"{API_SERVER_URL}/query/document-search", - json=search_request_body, - headers=test_user_1.headers, - ) - result.raise_for_status() - found_docs = result.json()["top_documents"] - danswer_doc_message_strings = [doc["content"] for doc in found_docs] print( "\ntop_documents content after prune for test_user_1: ", danswer_doc_message_strings,