Tenant integration tests

.github/workflows/pr-Integration-tests.yml

@@ -72,7 +72,7 @@ jobs:
           load: true
           cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/model-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }}
           cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/model-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max
-
+      
       - name: Build integration test Docker image
         uses: ./.github/actions/custom-build-and-push
         with:
@@ -85,7 +85,58 @@ jobs:
           cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/integration/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }}
           cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/integration/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max
 
-      - name: Start Docker containers
+      # Start containers for multi-tenant tests
+      - name: Start Docker containers for multi-tenant tests
+        run: |
+          cd deployment/docker_compose
+          ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true \
+          MULTI_TENANT=true \
+          AUTH_TYPE=basic \
+          REQUIRE_EMAIL_VERIFICATION=false \
+          DISABLE_TELEMETRY=true \
+          IMAGE_TAG=test \
+          docker compose -f docker-compose.dev.yml -p danswer-stack up -d
+        id: start_docker_multi_tenant
+
+      # In practice, `cloud` Auth type would require OAUTH credentials to be set.
+      - name: Run Multi-Tenant Integration Tests
+        run: |
+          echo "Running integration tests..."
+          docker run --rm --network danswer-stack_default \
+            --name test-runner \
+            -e POSTGRES_HOST=relational_db \
+            -e POSTGRES_USER=postgres \
+            -e POSTGRES_PASSWORD=password \
+            -e POSTGRES_DB=postgres \
+            -e VESPA_HOST=index \
+            -e REDIS_HOST=cache \
+            -e API_SERVER_HOST=api_server \
+            -e OPENAI_API_KEY=${OPENAI_API_KEY} \
+            -e SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN} \
+            -e TEST_WEB_HOSTNAME=test-runner \
+            -e AUTH_TYPE=cloud \
+            -e MULTI_TENANT=true \
+            danswer/danswer-integration:test \
+            /app/tests/integration/multitenant_tests
+        continue-on-error: true
+        id: run_multitenant_tests
+
+      - name: Check multi-tenant test results
+        run: |
+          if [ ${{ steps.run_tests.outcome }} == 'failure' ]; then
+            echo "Integration tests failed. Exiting with error."
+            exit 1
+          else
+            echo "All integration tests passed successfully."
+          fi 
+
+      - name: Stop multi-tenant Docker containers
+        run: |
+          cd deployment/docker_compose
+          docker compose -f docker-compose.dev.yml -p danswer-stack down -v
+
+
+      - name: Start Docker containers 
         run: |
           cd deployment/docker_compose
           ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true \
@@ -94,7 +145,7 @@ jobs:
           DISABLE_TELEMETRY=true \
           IMAGE_TAG=test \
           docker compose -f docker-compose.dev.yml -p danswer-stack up -d
-        id: start_docker
+        id: start_dockerstop m
 
       - name: Wait for service to be ready
         run: |
@@ -130,7 +181,7 @@ jobs:
           done
           echo "Finished waiting for service."
 
-      - name: Run integration tests
+      - name: Run Standard Integration Tests
         run: |
           echo "Running integration tests..."
           docker run --rm --network danswer-stack_default \
@@ -145,7 +196,8 @@ jobs:
             -e OPENAI_API_KEY=${OPENAI_API_KEY} \
             -e SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN} \
             -e TEST_WEB_HOSTNAME=test-runner \
-            danswer/danswer-integration:test
+            danswer/danswer-integration:test \
+            /app/tests/integration/tests
         continue-on-error: true
         id: run_tests
 
@@ -158,6 +210,11 @@ jobs:
             echo "All integration tests passed successfully."
           fi
 
+      - name: Stop Docker containers
+        run: |
+          cd deployment/docker_compose
+          docker compose -f docker-compose.dev.yml -p danswer-stack down -v
+
       - name: Save Docker logs
         if: success() || failure()
         run: |

backend/danswer/auth/users.py

@@ -58,6 +58,7 @@
 from danswer.auth.schemas import UserUpdate
 from danswer.configs.app_configs import AUTH_TYPE
 from danswer.configs.app_configs import DISABLE_AUTH
+from danswer.configs.app_configs import DISABLE_VERIFICATION
 from danswer.configs.app_configs import EMAIL_FROM
 from danswer.configs.app_configs import MULTI_TENANT
 from danswer.configs.app_configs import REQUIRE_EMAIL_VERIFICATION
@@ -132,7 +133,9 @@ def get_display_email(email: str | None, space_less: bool = False) -> str:
 def user_needs_to_be_verified() -> bool:
     # all other auth types besides basic should require users to be
     # verified
-    return AUTH_TYPE != AuthType.BASIC or REQUIRE_EMAIL_VERIFICATION
+    return not DISABLE_VERIFICATION and (
+        AUTH_TYPE != AuthType.BASIC or REQUIRE_EMAIL_VERIFICATION
+    )
 
 
 def verify_email_is_invited(email: str) -> None:

backend/danswer/configs/app_configs.py

@@ -43,6 +43,9 @@
 AUTH_TYPE = AuthType((os.environ.get("AUTH_TYPE") or AuthType.DISABLED.value).lower())
 DISABLE_AUTH = AUTH_TYPE == AuthType.DISABLED
 
+# Necessary for cloud integration tests
+DISABLE_VERIFICATION = os.environ.get("DISABLE_VERIFICATION", "").lower() == "true"
+
 # Encryption key secret is used to encrypt connector credentials, api keys, and other sensitive
 # information. This provides an extra layer of security on top of Postgres access controls
 # and is available in Danswer EE

backend/danswer/server/danswer_api/ingestion.py

@@ -9,6 +9,7 @@
 from danswer.db.connector_credential_pair import get_connector_credential_pair_from_id
 from danswer.db.document import get_documents_by_cc_pair
 from danswer.db.document import get_ingestion_documents
+from danswer.db.engine import get_current_tenant_id
 from danswer.db.engine import get_session
 from danswer.db.models import User
 from danswer.db.search_settings import get_current_search_settings
@@ -67,6 +68,7 @@ def upsert_ingestion_doc(
     doc_info: IngestionDocument,
     _: User | None = Depends(api_key_dep),
     db_session: Session = Depends(get_session),
+    tenant_id: str = Depends(get_current_tenant_id),
 ) -> IngestionResult:
     doc_info.document.from_ingestion_api = True
 
@@ -101,6 +103,7 @@ def upsert_ingestion_doc(
         document_index=curr_doc_index,
         ignore_time_skip=True,
         db_session=db_session,
+        tenant_id=tenant_id,
     )
 
     new_doc, __chunk_count = indexing_pipeline(
@@ -134,6 +137,7 @@ def upsert_ingestion_doc(
             document_index=sec_doc_index,
             ignore_time_skip=True,
             db_session=db_session,
+            tenant_id=tenant_id,
         )
 
         sec_ind_pipeline(

backend/ee/danswer/server/middleware/tenant_tracking.py

@@ -22,7 +22,6 @@ async def set_tenant_id(
     ) -> Response:
         try:
             logger.info(f"Request route: {request.url.path}")
-
             if not MULTI_TENANT:
                 tenant_id = POSTGRES_DEFAULT_SCHEMA
             else:

backend/tests/integration/Dockerfile

@@ -83,4 +83,5 @@ COPY ./tests/integration /app/tests/integration
 
 ENV PYTHONPATH=/app
 
-CMD ["pytest", "-s", "/app/tests/integration"]
+ENTRYPOINT ["pytest", "-s"]
+CMD ["/app/tests/integration", "--ignore=/app/tests/integration/multitenant_tests"]

backend/tests/integration/common_utils/managers/chat.py

@@ -23,7 +23,7 @@
 class ChatSessionManager:
     @staticmethod
     def create(
-        persona_id: int = -1,
+        persona_id: int = 0,
         description: str = "Test chat session",
         user_performing_action: DATestUser | None = None,
     ) -> DATestChatSession:

backend/tests/integration/common_utils/managers/credential.py

@@ -32,6 +32,7 @@ def create(
             "curator_public": curator_public,
             "groups": groups or [],
         }
+
         response = requests.post(
             url=f"{API_SERVER_URL}/manage/credential",
             json=credential_request,

backend/tests/integration/common_utils/managers/tenant.py

@@ -0,0 +1,82 @@
+from datetime import datetime
+from datetime import timedelta
+
+import jwt
+import requests
+
+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
+
+
+def generate_auth_token() -> str:
+    payload = {
+        "iss": "control_plane",
+        "exp": datetime.utcnow() + timedelta(minutes=5),
+        "iat": datetime.utcnow(),
+        "scope": "tenant:create",
+    }
+    token = jwt.encode(payload, "", algorithm="HS256")
+    return token
+
+
+class TenantManager:
+    @staticmethod
+    def create(
+        tenant_id: str | None = None,
+        initial_admin_email: str | None = None,
+    ) -> dict[str, str]:
+        body = {
+            "tenant_id": tenant_id,
+            "initial_admin_email": initial_admin_email,
+        }
+
+        token = generate_auth_token()
+        headers = {
+            "Authorization": f"Bearer {token}",
+            "X-API-KEY": "",
+            "Content-Type": "application/json",
+        }
+
+        response = requests.post(
+            url=f"{API_SERVER_URL}/tenants/create",
+            json=body,
+            headers=headers,
+        )
+
+        response.raise_for_status()
+
+        return response.json()
+
+    @staticmethod
+    def get_all_users(
+        user_performing_action: DATestUser | None = None,
+    ) -> AllUsersResponse:
+        response = requests.get(
+            url=f"{API_SERVER_URL}/manage/users",
+            headers=user_performing_action.headers
+            if user_performing_action
+            else GENERAL_HEADERS,
+        )
+        response.raise_for_status()
+
+        data = response.json()
+        return 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"],
+        )
+
+    @staticmethod
+    def verify_user_in_tenant(
+        user: DATestUser, user_performing_action: DATestUser | None = None
+    ) -> None:
+        all_users = TenantManager.get_all_users(user_performing_action)
+        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 in tenant")

backend/tests/integration/common_utils/managers/user.py

@@ -65,15 +65,23 @@ def login_as_user(test_user: DATestUser) -> DATestUser:
             data=data,
             headers=headers,
         )
+
         response.raise_for_status()
-        result_cookie = next(iter(response.cookies), None)
 
-        if not result_cookie:
+        cookies = response.cookies.get_dict()
+        session_cookie = cookies.get("fastapiusersauth")
+        tenant_details_cookie = cookies.get("tenant_details")
+
+        if not session_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
+
+        # Set both cookies in the headers
+        test_user.headers["Cookie"] = (
+            f"fastapiusersauth={session_cookie}; "
+            f"tenant_details={tenant_details_cookie}"
+        )
         return test_user
 
     @staticmethod