[PLT-1205] Add additional UG Tests / Proper Error Handling (#1712)

Author: adrian-chang

libs/labelbox/src/labelbox/exceptions.py

@@ -32,17 +32,21 @@ class AuthorizationError(LabelboxError):
 class ResourceNotFoundError(LabelboxError):
     """Exception raised when a given resource is not found. """
 
-    def __init__(self, db_object_type, params):
-        """ Constructor.
+    def __init__(self, db_object_type=None, params=None, message=None):
+        """Constructor for the ResourceNotFoundException class.
 
         Args:
-            db_object_type (type): A labelbox.schema.DbObject subtype.
-            params (dict): Dict of params identifying the sought resource.
+            db_object_type (type): A subtype of labelbox.schema.DbObject.
+            params (dict): A dictionary of parameters identifying the sought resource.
+            message (str): An optional message to include in the exception.
         """
-        super().__init__("Resource '%s' not found for params: %r" %
-                         (db_object_type.type_name(), params))
-        self.db_object_type = db_object_type
-        self.params = params
+        if message is not None:
+            super().__init__(message)
+        else:
+            super().__init__("Resource '%s' not found for params: %r" %
+                                (db_object_type.type_name(), params))
+            self.db_object_type = db_object_type
+            self.params = params
 
 
 class ResourceConflict(LabelboxError):

libs/labelbox/src/labelbox/schema/user_group.py

@@ -7,7 +7,7 @@
 from labelbox.pydantic_compat import BaseModel
 from labelbox.schema.user import User
 from labelbox.schema.project import Project
-from labelbox.exceptions import UnprocessableEntityError, InvalidQueryError
+from labelbox.exceptions import UnprocessableEntityError, MalformedQueryException, ResourceNotFoundError
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.ontology_kind import EditorTaskType
 from labelbox.schema.media_type import MediaType
@@ -105,14 +105,11 @@ def get(self) -> "UserGroup":
         about the user group, including its name, color, projects, and members. The fetched
         information is then used to update the corresponding attributes of the `Group` object.
 
-        Args:
-            id (str): The ID of the user group to fetch.
-
         Returns:
-            UserGroup of passed in ID (self)
+            UserGroup: The updated `UserGroup` object.
 
         Raises:
-            InvalidQueryError: If the query fails to fetch the group information.
+            ResourceNotFoundError: If the query fails to fetch the group information.
         """
         query = """
             query GetUserGroupPyApi($id: ID!) {
@@ -142,7 +139,7 @@ def get(self) -> "UserGroup":
         }
         result = self.client.execute(query, params)
         if not result:
-            raise InvalidQueryError("Failed to fetch group, ID likely incorrect")
+            raise ResourceNotFoundError(message="Failed to get user group as user group does not exist")
         self.name = result["userGroup"]["name"]
         self.color = UserGroupColor(result["userGroup"]["color"])
         self.projects = self._get_projects_set(result["userGroup"]["projects"]["nodes"])
@@ -157,7 +154,8 @@ def update(self) -> "UserGroup":
             UserGroup: The updated UserGroup object. (self)
 
         Raises:
-            UnprocessableEntityError: If the update fails.
+            ResourceNotFoundError: If the update fails due to unknown user group
+            UnprocessableEntityError: If the update fails due to a malformed input
         """
         query = """
         mutation UpdateUserGroupPyApi($id: ID!, $name: String!, $color: String!, $projectIds: [String!]!, $userIds: [String!]!) {
@@ -199,9 +197,12 @@ def update(self) -> "UserGroup":
                 user.uid for user in self.users
             ]
         }
-        result = self.client.execute(query, params)
-        if not result:
-            raise UnprocessableEntityError("Failed to update user group, either user group name is in use currently, or provided user or projects don't exist")
+        try:
+            result = self.client.execute(query, params)
+            if not result:
+                raise ResourceNotFoundError(message="Failed to update user group as user group does not exist")
+        except MalformedQueryException as e:
+            raise UnprocessableEntityError("Failed to update user group") from e
         return self
 
     def create(self) -> "UserGroup":
@@ -261,9 +262,15 @@ def create(self) -> "UserGroup":
                 user.uid for user in self.users
             ]
         }
-        result = self.client.execute(query, params)
-        if not result:
-            raise ResourceCreationError("Failed to create user group, either user group name is in use currently, or provided user or projects don't exist")
+        result = None
+        error = None
+        try: 
+            result = self.client.execute(query, params)
+        except Exception as e:
+            error = e
+        if not result or error:
+            # this is client side only, server doesn't have an equivalent error
+            raise ResourceCreationError(f"Failed to create user group, either user group name is in use currently, or provided user or projects don't exist server error: {error}")
         result = result["createUserGroup"]["group"]
         self.id = result["id"]
         return self
@@ -280,7 +287,7 @@ def delete(self) -> bool:
             bool: True if the user group was successfully deleted, False otherwise.
         
         Raises:
-            UnprocessableEntityError: If the deletion of the user group fails.
+            ResourceNotFoundError: If the deletion of the user group fails due to not existing
         """
         query = """
         mutation DeleteUserGroupPyApi($id: ID!) {
@@ -292,7 +299,7 @@ def delete(self) -> bool:
         params = {"id": self.id}
         result = self.client.execute(query, params)
         if not result:
-            raise UnprocessableEntityError("Failed to delete user group")
+            raise ResourceNotFoundError(message="Failed to delete user group as user group does not exist")
         return result["deleteUserGroup"]["success"]
 
     def get_user_groups(self) -> Iterator["UserGroup"]:
@@ -306,8 +313,8 @@ def get_user_groups(self) -> Iterator["UserGroup"]:
             Iterator[UserGroup]: An iterator over the user groups.
         """
         query = """
-            query GetUserGroupsPyApi {
-                userGroups {
+            query GetUserGroupsPyApi($after: String) {
+                userGroups(after: $after) {
                     nodes {
                         id
                         name
@@ -334,7 +341,7 @@ def get_user_groups(self) -> Iterator["UserGroup"]:
         nextCursor = None
         while True:
             userGroups = self.client.execute(
-                query, {"nextCursor": nextCursor})["userGroups"]
+                query, {"after": nextCursor})["userGroups"]
             if not userGroups:
                 return
                 yield
@@ -348,9 +355,9 @@ def get_user_groups(self) -> Iterator["UserGroup"]:
                 userGroup.projects = self._get_projects_set(group["projects"]["nodes"])
                 yield userGroup
             nextCursor = userGroups["nextCursor"]
-            # this doesn't seem to be implemented right now to return a value other than null from the api
             if not nextCursor:
-                break
+                return
+                yield
 
     def _get_users_set(self, user_nodes):
         """

libs/labelbox/tests/conftest.py

@@ -91,7 +91,9 @@ def ephemeral_endpoint() -> str:
 
 
 def graphql_url(environ: str) -> str:
-    if environ == Environ.PROD:
+    if environ == Environ.LOCAL:
+        return 'http://localhost:3000/api/graphql'
+    elif environ == Environ.PROD:
         return 'https://api.labelbox.com/graphql'
     elif environ == Environ.STAGING:
         return 'https://api.lb-stage.xyz/graphql'
@@ -107,7 +109,9 @@ def graphql_url(environ: str) -> str:
 
 
 def rest_url(environ: str) -> str:
-    if environ == Environ.PROD:
+    if environ == Environ.LOCAL:
+        return 'http://localhost:3000/api/v1'
+    elif environ == Environ.PROD:
         return 'https://api.labelbox.com/api/v1'
     elif environ == Environ.STAGING:
         return 'https://api.lb-stage.xyz/api/v1'
@@ -124,7 +128,8 @@ def rest_url(environ: str) -> str:
 def testing_api_key(environ: Environ) -> str:
     keys = [
         f"LABELBOX_TEST_API_KEY_{environ.value.upper()}",
-        "LABELBOX_TEST_API_KEY"
+        "LABELBOX_TEST_API_KEY",
+        "LABELBOX_API_KEY"
     ]
     for key in keys:
         value = os.environ.get(key)
@@ -311,10 +316,16 @@ def environ() -> Environ:
     'prod' or 'staging'
     Make sure to set LABELBOX_TEST_ENVIRON in .github/workflows/python-package.yaml
     """
-    try:
-        return Environ(os.environ['LABELBOX_TEST_ENVIRON'])
-    except KeyError:
-        raise Exception(f'Missing LABELBOX_TEST_ENVIRON in: {os.environ}')
+    keys = [
+        "LABELBOX_TEST_ENV",
+        "LABELBOX_TEST_ENVIRON",
+        "LABELBOX_ENV"
+    ]
+    for key in keys:
+        value = os.environ.get(key)
+        if value is not None:
+            return Environ(value)
+    raise Exception(f'Missing env key in: {os.environ}')
 
 
 def cancel_invite(client, invite_id):

libs/labelbox/tests/integration/schema/test_user_group.py

@@ -1,7 +1,9 @@
 import pytest
 import faker
+from uuid import uuid4
 from labelbox import Client
 from labelbox.schema.user_group import UserGroup, UserGroupColor
+from labelbox.exceptions import ResourceNotFoundError, ResourceCreationError, UnprocessableEntityError
 
 data = faker.Faker()
 
@@ -29,47 +31,132 @@ def test_existing_user_groups(user_group, client):
     assert user_group.color == user_group_equal.color
 
 
+def test_cannot_get_user_group_with_invalid_id(client):
+    user_group = UserGroup(client=client, id=str(uuid4()))
+    with pytest.raises(ResourceNotFoundError):
+        user_group.get()
+
+
+def test_throw_error_when_retrieving_deleted_group(client):
+    user_group = UserGroup(client=client, name=data.name())
+    user_group.create()
+
+    assert user_group.get() is not None
+    user_group.delete()
+
+    with pytest.raises(ResourceNotFoundError):
+        user_group.get()
+
+
+def test_create_user_group_no_name(client):
+    # Create a new user group
+    with pytest.raises(ResourceCreationError):
+        user_group = UserGroup(client)
+        user_group.name = "   "
+        user_group.color = UserGroupColor.BLUE
+        user_group.create()
+
+
+def test_cannot_create_group_with_same_name(client, user_group):
+    with pytest.raises(ResourceCreationError):
+        user_group_2 = UserGroup(client=client, name=user_group.name)
+        user_group_2.create()
+
+
 def test_create_user_group(user_group):
     # Verify that the user group was created successfully
     assert user_group.id is not None
     assert user_group.name is not None
     assert user_group.color == UserGroupColor.BLUE
 
 
+def test_create_user_group_advanced(client, project_pack):
+    group_name = data.name()
+    # Create a new user group
+    user_group = UserGroup(client)
+    user_group.name = group_name
+    user_group.color = UserGroupColor.BLUE
+    users = list(client.get_users())
+    projects = project_pack
+    user = users[0]
+    project = projects[0]
+    user_group.users.add(user)
+    user_group.projects.add(project)
+
+    user_group.create()
+
+    assert user_group.id is not None
+    assert user_group.name is not None
+    assert user_group.color == UserGroupColor.BLUE
+    assert project in user_group.projects
+    assert user in user_group.users
+
+    user_group.delete()
+
+
 def test_update_user_group(user_group):
     # Update the user group
     group_name = data.name()
     user_group.name = group_name
     user_group.color = UserGroupColor.PURPLE
-    user_group.update()
+    updated_user_group = user_group.update()
 
     # Verify that the user group was updated successfully
+    assert user_group.name == updated_user_group.name
     assert user_group.name == group_name
+    assert user_group.color == updated_user_group.color
     assert user_group.color == UserGroupColor.PURPLE
 
 
-def test_get_user_groups(user_group, client):
-    # Get all user groups
-    user_groups_old = list(UserGroup(client).get_user_groups())
+def test_cannot_update_name_to_empty_string(user_group):
+    with pytest.raises(UnprocessableEntityError):
+        user_group.name = ""
+        user_group.update()
 
-    # manual delete for iterators
-    group_name = data.name()
-    user_group = UserGroup(client)
-    user_group.name = group_name
-    user_group.create()
 
-    user_groups_new = list(UserGroup(client).get_user_groups())
+def test_cannot_update_group_id(user_group):
+    old_id = user_group.id
+    with pytest.raises(ResourceNotFoundError):
+        user_group.id = str(uuid4())
+        user_group.update()
+    # prevent leak
+    user_group.id = old_id
 
-    # Verify that at least one user group is returned
-    assert len(user_groups_new) > 0
-    assert len(user_groups_new) == len(user_groups_old) + 1
 
-    # Verify that each user group has a valid ID and name
-    for user_group in user_groups_new:
-        assert user_group.id is not None
-        assert user_group.name is not None
+def test_get_user_groups_with_creation_deletion(client):
+    user_group = None
+    try: 
+        # Get all user groups
+        user_groups = list(UserGroup(client).get_user_groups())
 
-    user_group.delete()
+        # manual delete for iterators
+        group_name = data.name()
+        user_group = UserGroup(client)
+        user_group.name = group_name
+        user_group.create()
+
+        user_groups_post_creation = list(UserGroup(client).get_user_groups())
+
+        # Verify that at least one user group is returned
+        assert len(user_groups_post_creation) > 0
+        assert len(user_groups_post_creation) == len(user_groups) + 1
+
+        # Verify that each user group has a valid ID and name
+        for user_group in user_groups_post_creation:
+            assert user_group.id is not None
+            assert user_group.name is not None
+
+        user_group.delete()
+        user_group = None
+
+        user_groups_post_deletion = list(UserGroup(client).get_user_groups())
+
+        assert len(user_groups_post_deletion) > 0
+        assert len(user_groups_post_deletion) == len(user_groups_post_creation) - 1
+
+    finally:
+        if user_group:
+            user_group.delete()
 
 
 # project_pack creates two projects
@@ -89,6 +176,22 @@ def test_update_user_group(user_group, client, project_pack):
     assert project in user_group.projects
 
 
+def test_delete_user_group_with_same_id(client):
+    user_group_1 = UserGroup(client, name=data.name())
+    user_group_1.create()
+    user_group_1.delete()
+    user_group_2 = UserGroup(client=client, id=user_group_1.id)
+
+    with pytest.raises(ResourceNotFoundError):
+        user_group_2.delete()
+
+
+def test_throw_error_when_deleting_invalid_id_group(client):
+    with pytest.raises(ResourceNotFoundError):
+        user_group = UserGroup(client=client, id=str(uuid4()))
+        user_group.delete()
+
+
 if __name__ == "__main__":
     import subprocess
     subprocess.call(["pytest", "-v", __file__])