Add multi-tenant user invitation flow test (#4360)

backend/ee/onyx/server/tenants/user_mapping.py

@@ -70,6 +70,7 @@ def add_users_to_tenant(emails: list[str], tenant_id: str) -> None:
     """
     Add users to a tenant with proper transaction handling.
     Checks if users already have a tenant mapping to avoid duplicates.
+    If a user already has an active mapping to any tenant, the new mapping will be added as inactive.
     """
     with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db_session:
         try:
@@ -88,9 +89,25 @@ def add_users_to_tenant(emails: list[str], tenant_id: str) -> None:
                     .first()
                 )
 
+                # If user already has an active mapping, add this one as inactive
                 if not existing_mapping:
-                    # Only add if mapping doesn't exist
-                    db_session.add(UserTenantMapping(email=email, tenant_id=tenant_id))
+                    # Check if the user already has an active mapping to any tenant
+                    has_active_mapping = (
+                        db_session.query(UserTenantMapping)
+                        .filter(
+                            UserTenantMapping.email == email,
+                            UserTenantMapping.active == True,  # noqa: E712
+                        )
+                        .first()
+                    )
+
+                    db_session.add(
+                        UserTenantMapping(
+                            email=email,
+                            tenant_id=tenant_id,
+                            active=False if has_active_mapping else True,
+                        )
+                    )
 
             # Commit the transaction
             db_session.commit()

backend/onyx/server/manage/users.py

@@ -313,7 +313,7 @@ def bulk_invite_users(
             detail=f"Invalid email address: {email} - {str(e)}",
         )
 
-    if MULTI_TENANT and not DEV_MODE:
+    if MULTI_TENANT:
         try:
             fetch_ee_implementation_or_noop(
                 "onyx.server.tenants.provisioning", "add_users_to_tenant", None
@@ -335,7 +335,7 @@ def bulk_invite_users(
         except Exception as e:
             logger.error(f"Error sending email invite to invited users: {e}")
 
-    if not MULTI_TENANT:
+    if not MULTI_TENANT or DEV_MODE:
         return number_of_invited_users
 
     # for billing purposes, write to the control plane about the number of new users
@@ -376,7 +376,7 @@ def remove_invited_user(
     number_of_invited_users = write_invited_users(remaining_users)
 
     try:
-        if MULTI_TENANT:
+        if MULTI_TENANT and not DEV_MODE:
             fetch_ee_implementation_or_noop(
                 "onyx.server.tenants.billing", "register_tenant_users", None
             )(tenant_id, get_total_users_count(db_session))

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

@@ -9,7 +9,9 @@ from requests import HTTPError
 from onyx.auth.schemas import UserRole
 from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
 from onyx.server.documents.models import PaginatedReturn
+from onyx.server.manage.models import UserInfo
 from onyx.server.models import FullUserSnapshot
+from onyx.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
@@ -245,3 +247,69 @@ class UserManager:
             total_items=data["total_items"],
         )
         return paginated_result
+
+    @staticmethod
+    def invite_user(
+        user_to_invite_email: str, user_performing_action: DATestUser
+    ) -> None:
+        """Invite a user by email to join the organization.
+
+        Args:
+            user_to_invite_email: Email of the user to invite
+            user_performing_action: User with admin permissions performing the invitation
+        """
+        response = requests.put(
+            url=f"{API_SERVER_URL}/manage/admin/users",
+            headers=user_performing_action.headers,
+            json={"emails": [user_to_invite_email]},
+        )
+        response.raise_for_status()
+
+    @staticmethod
+    def accept_invitation(tenant_id: str, user_performing_action: DATestUser) -> None:
+        """Accept an invitation to join the organization.
+
+        Args:
+            tenant_id: ID of the tenant/organization to accept invitation for
+            user_performing_action: User accepting the invitation
+        """
+        response = requests.post(
+            url=f"{API_SERVER_URL}/tenants/users/invite/accept",
+            headers=user_performing_action.headers,
+            json={"tenant_id": tenant_id},
+        )
+        response.raise_for_status()
+
+    @staticmethod
+    def get_invited_users(
+        user_performing_action: DATestUser,
+    ) -> list[InvitedUserSnapshot]:
+        """Get a list of all invited users.
+
+        Args:
+            user_performing_action: User with admin permissions performing the action
+
+        Returns:
+            List of invited user snapshots
+        """
+        response = requests.get(
+            url=f"{API_SERVER_URL}/manage/users/invited",
+            headers=user_performing_action.headers,
+        )
+        response.raise_for_status()
+
+        return [InvitedUserSnapshot(**user) for user in response.json()]
+
+    @staticmethod
+    def get_user_info(user_performing_action: DATestUser) -> UserInfo:
+        """Get user info for the current user.
+
+        Args:
+            user_performing_action: User performing the action
+        """
+        response = requests.get(
+            url=f"{API_SERVER_URL}/me",
+            headers=user_performing_action.headers,
+        )
+        response.raise_for_status()
+        return UserInfo(**response.json())

backend/tests/integration/multitenant_tests/invitation/invite_various_organizations.py

@@ -0,0 +1,70 @@
+from onyx.db.models import UserRole
+from tests.integration.common_utils.managers.user import UserManager
+from tests.integration.common_utils.test_models import DATestUser
+
+INVITED_BASIC_USER = "basic_user"
+INVITED_BASIC_USER_EMAIL = "basic_user@test.com"
+
+
+def test_user_invitation_flow(reset_multitenant: None) -> None:
+    # Create first user (admin)
+    admin_user: DATestUser = UserManager.create(name="admin")
+    assert UserManager.is_role(admin_user, UserRole.ADMIN)
+
+    # Create second user
+    invited_user: DATestUser = UserManager.create(name="admin_invited")
+    assert UserManager.is_role(invited_user, UserRole.ADMIN)
+
+    # Admin user invites the previously registered and non-registered user
+    UserManager.invite_user(invited_user.email, admin_user)
+    UserManager.invite_user(INVITED_BASIC_USER_EMAIL, admin_user)
+
+    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"
+
+    # Get user info to check tenant information
+    user_info = UserManager.get_user_info(invited_user)
+
+    # Extract the tenant_id from the invitation
+    invited_tenant_id = (
+        user_info.tenant_info.invitation.tenant_id
+        if user_info.tenant_info and user_info.tenant_info.invitation
+        else None
+    )
+    assert invited_tenant_id is not None, "Expected to find an invitation tenant_id"
+
+    UserManager.accept_invitation(invited_tenant_id, invited_user)
+
+    # Get updated user info after accepting invitation
+    updated_user_info = UserManager.get_user_info(invited_user)
+
+    # 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"
+
+    # Verify the user has BASIC role in the organization
+    assert (
+        updated_user_info.role == UserRole.BASIC
+    ), f"Expected user to have BASIC role, but got {updated_user_info.role}"
+
+    # Verify user is in the organization
+    user_page = UserManager.get_user_page(
+        user_performing_action=admin_user, role_filter=[UserRole.BASIC]
+    )
+
+    # 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 "
+        f"in the organization. Available users: {invited_user_emails}"
+    )