Prevent workspace-wide users from being user group members

paulnoirel · paulnoirel · commit 3866a9790fb9 · 2025-07-02T23:48:14.000+01:00
diff --git a/libs/labelbox/src/labelbox/schema/user_group.py b/libs/labelbox/src/labelbox/schema/user_group.py
@@ -230,7 +230,7 @@ def update(self) -> UserGroup:
         Raises:
             ValueError: If group ID or name is not set, or if projects don't exist.
             ResourceNotFoundError: If the group or projects are not found.
-            UnprocessableEntityError: If user validation fails.
+            UnprocessableEntityError: If user validation fails or users have workspace-level org roles.
         """
         if not self.id:
             raise ValueError("Group id is required")
@@ -247,8 +247,11 @@ def update(self) -> UserGroup:
                 )
 
         # Filter eligible users and build user roles
-        eligible_users = self._filter_project_based_users()
-        user_roles = self._build_user_roles(eligible_users)
+        try:
+            eligible_users = self._filter_project_based_users()
+            user_roles = self._build_user_roles(eligible_users)
+        except ValueError as e:
+            raise UnprocessableEntityError(str(e)) from e
 
         query = """
         mutation UpdateUserGroupPyApi($id: ID!, $name: String!, $description: String, $color: String!, $projectIds: [ID!]!, $userRoles: [UserRoleInput!], $notifyMembers: Boolean) {
@@ -312,7 +315,7 @@ def create(self) -> UserGroup:
 
         Raises:
             ValueError: If group already has ID, name is invalid, or projects don't exist.
-            ResourceCreationError: If creation fails or user validation fails.
+            ResourceCreationError: If creation fails, user validation fails, or users have workspace-level org roles.
             ResourceConflict: If a group with the same name already exists.
         """
         if self.id:
@@ -330,8 +333,11 @@ def create(self) -> UserGroup:
                 )
 
         # Filter eligible users and build user roles
-        eligible_users = self._filter_project_based_users()
-        user_roles = self._build_user_roles(eligible_users)
+        try:
+            eligible_users = self._filter_project_based_users()
+            user_roles = self._build_user_roles(eligible_users)
+        except ValueError as e:
+            raise ResourceCreationError(str(e)) from e
 
         query = """
         mutation CreateUserGroupPyApi($name: String!, $description: String, $color: String!, $projectIds: [ID!]!, $userRoles: [UserRoleInput!]!, $notifyMembers: Boolean, $roleId: String, $searchQuery: AlignerrSearchServiceQuery) {
@@ -496,11 +502,14 @@ def get_user_groups(
     def _filter_project_based_users(self) -> Set[User]:
         """Filter users to only include users eligible for UserGroups.
 
-        Filters out users with specific admin organization roles that cannot be
-        added to UserGroups. Most users should be eligible.
+        Only project-based users (org role "NONE") can be added to UserGroups.
+        Users with any workspace-level organization role cannot be added.
 
         Returns:
             Set of users that are eligible to be added to the group.
+
+        Raises:
+            ValueError: If any user has a workspace-level organization role.
         """
         all_users = set()
         for member in self.members:
@@ -509,63 +518,52 @@ def _filter_project_based_users(self) -> Set[User]:
         if not all_users:
             return set()
 
-        user_ids = [user.uid for user in all_users]
-        query = """
-        query CheckUserOrgRolesPyApi($userIds: [ID!]!) {
-            users(where: {id_in: $userIds}) {
-                id
-                orgRole { id name }
-            }
-        }
-        """
+        # Check each user's org role directly
+        invalid_users = []
+        eligible_users = set()
 
-        try:
-            result = self.client.execute(query, {"userIds": user_ids})
-            if not result or "users" not in result:
-                return all_users  # Fallback: let server handle validation
-
-            # Check for users with org roles that cannot be used in UserGroups
-            # Only users with no org role (project-based users) can be assigned to UserGroups
-            eligible_user_ids = set()
-            invalid_users = []
-
-            for user_data in result["users"]:
-                org_role = user_data.get("orgRole")
-                user_id = user_data["id"]
-                user_email = user_data.get("email", "unknown")
-
-                if org_role is None:
-                    # Users with no org role (project-based users) are eligible
-                    eligible_user_ids.add(user_id)
+        for user in all_users:
+            try:
+                # Get the user's organization role directly
+                org_role = user.org_role()
+                if org_role is None or org_role.name.upper() == "NONE":
+                    # Users with no org role or "NONE" role are project-based and eligible
+                    eligible_users.add(user)
                 else:
-                    # Users with ANY workspace org role cannot be assigned to UserGroups
+                    # Users with any workspace org role cannot be assigned to UserGroups
                     invalid_users.append(
                         {
-                            "id": user_id,
-                            "email": user_email,
-                            "org_role": org_role.get("name"),
+                            "id": user.uid,
+                            "email": getattr(user, "email", "unknown"),
+                            "org_role": org_role.name,
                         }
                     )
+            except Exception as e:
+                # If we can't determine the user's role, treat as invalid for safety
+                invalid_users.append(
+                    {
+                        "id": user.uid,
+                        "email": getattr(user, "email", "unknown"),
+                        "org_role": f"unknown (error: {str(e)})",
+                    }
+                )
 
-            # Raise error if any invalid users found
-            if invalid_users:
-                error_details = []
-                for user in invalid_users:
-                    error_details.append(
-                        f"User {user['id']} ({user['email']}) has org role '{user['org_role']}'"
-                    )
-
-                raise ValueError(
-                    f"Cannot create UserGroup with users who have organization roles. "
-                    f"Only project-based users (no org role) can be assigned to UserGroups.\n"
-                    f"Invalid users:\n"
-                    + "\n".join(f"  • {detail}" for detail in error_details)
+        # Raise error if any invalid users found
+        if invalid_users:
+            error_details = []
+            for user in invalid_users:
+                error_details.append(
+                    f"User {user['id']} ({user['email']}) has org role '{user['org_role']}'"
                 )
 
-            return {user for user in all_users if user.uid in eligible_user_ids}
+            raise ValueError(
+                f"Cannot create UserGroup with users who have organization roles. "
+                f"Only project-based users (no org role or role 'NONE') can be assigned to UserGroups.\n"
+                f"Invalid users:\n"
+                + "\n".join(f"  • {detail}" for detail in error_details)
+            )
 
-        except Exception:
-            return all_users  # Fallback: let server handle validation
+        return eligible_users
 
     def _build_user_roles(
         self, eligible_users: Set[User]
diff --git a/libs/labelbox/tests/unit/schema/test_user_group.py b/libs/labelbox/tests/unit/schema/test_user_group.py
@@ -41,7 +41,10 @@ def group_user():
     user_values["createdAt"] = "2023-01-01T00:00:00Z"
     user_values["isExternalUser"] = False
     user_values["isViewer"] = False
-    return User(MagicMock(Client), user_values)
+    user = User(MagicMock(Client), user_values)
+    # Mock org_role() to return None (project-based user)
+    user.org_role = MagicMock(return_value=None)
+    return user
 
 
 @pytest.fixture
@@ -237,16 +240,6 @@ def test_update(self, group_user, group_project, mock_role):
         self.client.get_project.return_value = group_project
 
         self.client.execute.side_effect = [
-            # Mock user roles query response
-            {
-                "users": [
-                    {
-                        "id": "user_id",
-                        "email": "test@example.com",
-                        "orgRole": None,  # Project-based user
-                    }
-                ]
-            },
             # Mock update mutation response
             {
                 "updateUserGroupV3": {
@@ -458,16 +451,6 @@ def test_create(self, group_user, group_project, mock_role):
         self.client.get_project.return_value = group_project
 
         self.client.execute.side_effect = [
-            # Mock user roles query response
-            {
-                "users": [
-                    {
-                        "id": "user_id",
-                        "email": "test@example.com",
-                        "orgRole": None,  # Project-based user
-                    }
-                ]
-            },
             # Mock create mutation response
             {
                 "createUserGroupV3": {
diff --git a/test_usergroup_fix.py b/test_usergroup_fix.py
@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+"""
+Test script to demonstrate the fixed UserGroup validation behavior.
+This script shows how workspace-based users are now properly rejected.
+"""
+
+from unittest.mock import MagicMock, Mock
+from collections import defaultdict
+
+# Import the fixed UserGroup classes
+from libs.labelbox.src.labelbox.schema.user_group import UserGroup, UserGroupMember, UserGroupColor
+from libs.labelbox.src.labelbox.schema.user import User
+from libs.labelbox.src.labelbox.schema.role import Role
+from lbox.exceptions import ResourceCreationError, UnprocessableEntityError
+
+
+def create_mock_user(user_id: str, email: str, org_role_name: str = None):
+    """Create a mock user with specified org role."""
+    client = MagicMock()
+    user_values = defaultdict(lambda: None)
+    user_values["id"] = user_id
+    user_values["email"] = email
+    
+    user = User(client, user_values)
+    
+    # Mock the org_role() method
+    if org_role_name is None or org_role_name.upper() == "NONE":
+        user.org_role = Mock(return_value=Mock(name="NONE"))
+    else:
+        user.org_role = Mock(return_value=Mock(name=org_role_name))
+    
+    return user
+
+
+def create_mock_role(role_name: str):
+    """Create a mock role."""
+    client = MagicMock()
+    role_values = defaultdict(lambda: None)
+    role_values["id"] = f"role_{role_name.lower()}"
+    role_values["name"] = role_name
+    return Role(client, role_values)
+
+
+def test_workspace_user_rejection():
+    """Test that workspace-based users are properly rejected."""
+    print("🧪 Testing UserGroup validation fixes...")
+    
+    client = MagicMock()
+    
+    # Create test users
+    project_user = create_mock_user("user1", "project@example.com", "NONE")  # Valid
+    workspace_labeler = create_mock_user("user2", "labeler@example.com", "LABELER")  # Invalid
+    workspace_admin = create_mock_user("user3", "admin@example.com", "ADMIN")  # Invalid
+    
+    # Create test role
+    labeler_role = create_mock_role("LABELER")
+    
+    print("\n✅ Test 1: Project-based user should be accepted")
+    try:
+        user_group = UserGroup(
+            client=client,
+            name="Valid Group",
+            color=UserGroupColor.BLUE,
+            members={UserGroupMember(user=project_user, role=labeler_role)}
+        )
+        # This should work - call the validation method directly
+        eligible_users = user_group._filter_project_based_users()
+        print(f"   ✓ Project-based user accepted: {len(eligible_users)} eligible users")
+    except Exception as e:
+        print(f"   ❌ Unexpected error: {e}")
+    
+    print("\n❌ Test 2: Workspace labeler should be rejected")
+    try:
+        user_group = UserGroup(
+            client=client,
+            name="Invalid Group",
+            color=UserGroupColor.BLUE,
+            members={UserGroupMember(user=workspace_labeler, role=labeler_role)}
+        )
+        # This should fail
+        eligible_users = user_group._filter_project_based_users()
+        print(f"   ❌ ERROR: Workspace labeler was incorrectly accepted!")
+    except ValueError as e:
+        print(f"   ✓ Workspace labeler correctly rejected: {str(e)[:100]}...")
+    
+    print("\n❌ Test 3: Workspace admin should be rejected")
+    try:
+        user_group = UserGroup(
+            client=client,
+            name="Invalid Group 2",
+            color=UserGroupColor.BLUE,
+            members={UserGroupMember(user=workspace_admin, role=labeler_role)}
+        )
+        # This should fail
+        eligible_users = user_group._filter_project_based_users()
+        print(f"   ❌ ERROR: Workspace admin was incorrectly accepted!")
+    except ValueError as e:
+        print(f"   ✓ Workspace admin correctly rejected: {str(e)[:100]}...")
+    
+    print("\n❌ Test 4: Mixed users - should reject all if any are invalid")
+    try:
+        user_group = UserGroup(
+            client=client,
+            name="Mixed Group",
+            color=UserGroupColor.BLUE,
+            members={
+                UserGroupMember(user=project_user, role=labeler_role),  # Valid
+                UserGroupMember(user=workspace_labeler, role=labeler_role),  # Invalid
+            }
+        )
+        # This should fail because of the workspace labeler
+        eligible_users = user_group._filter_project_based_users()
+        print(f"   ❌ ERROR: Mixed group with workspace user was incorrectly accepted!")
+    except ValueError as e:
+        print(f"   ✓ Mixed group correctly rejected: {str(e)[:100]}...")
+    
+    print("\n🎉 All tests completed!")
+    print("\n📋 Summary:")
+    print("   • Project-based users (org role 'NONE' or null) ✅ ACCEPTED")
+    print("   • Workspace-based users (any org role) ❌ REJECTED")
+    print("   • Clear error messages provided")
+    print("   • No more silent failures or server crashes")
+
+
+if __name__ == "__main__":
+    test_workspace_user_rejection() 
