address feedback 6

Author: Adrian Chang

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

@@ -338,8 +338,7 @@ def delete(self) -> bool:
             raise UnprocessableEntityError("Failed to delete user group")
         return result["deleteUserGroup"]["success"]
 
-    @staticmethod
-    def get_user_groups(client: Client) -> Iterator["UserGroup"]:
+    def get_user_groups(self) -> Iterator["UserGroup"]:
         """
         Gets all user groups in Labelbox.
 
@@ -377,14 +376,14 @@ def get_user_groups(client: Client) -> Iterator["UserGroup"]:
         """
         nextCursor = None
         while True:
-            userGroups = client.execute(
+            userGroups = self.client.execute(
                 query, {"nextCursor": nextCursor})["userGroups"]
             if not userGroups:
                 return
                 yield
             groups = userGroups["nodes"]
             for group in groups:
-                userGroup = UserGroup(client)
+                userGroup = UserGroup(self.client)
                 userGroup.id = group["id"]
                 userGroup.name = group["name"]
                 userGroup.color = UserGroupColor(group["color"])

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

@@ -50,15 +50,15 @@ def test_update_user_group(user_group):
 
 def test_get_user_groups(user_group, client):
     # Get all user groups
-    user_groups_old = list(UserGroup.get_user_groups(client))
+    user_groups_old = list(UserGroup(client).get_user_groups())
 
     # 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.get_user_groups(client))
+    user_groups_new = list(UserGroup(client).get_user_groups())
 
     # Verify that at least one user group is returned
     assert len(user_groups_new) > 0

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

@@ -283,7 +283,7 @@ def test_delete(self):
     def test_user_groups_empty(self):
         self.client.execute.return_value = {"userGroups": None}
 
-        user_groups = list(UserGroup.get_user_groups(self.client))
+        user_groups = list(UserGroup(self.client).get_user_groups())
 
         assert len(user_groups) == 0
 
@@ -362,8 +362,10 @@ def test_user_groups(self):
             }
         }
 
-        user_groups = list(UserGroup.get_user_groups(self.client))
+        user_groups = list(UserGroup(self.client).get_user_groups())
+        execute = self.client.execute.call_args[0]
 
+        assert "GetUserGroupsPyApi" in execute[0]
         assert len(user_groups) == 3
 
         # Check the attributes of the first user group