fixed test

Author: Matt Sokoloff

labelbox/schema/role.py

@@ -43,14 +43,11 @@ def __getitem__(self, name):
             )
 
     def __iter__(self):
+        self.key_iter = iter(self.valid_roles)
         return self
 
     def __next__(self):
-        try:
-            key = self.valid_roles[self.index]
-        except:
-            raise StopIteration
-        self.index += 1
+        key = next(self.key_iter)
         return getattr(self, key)
 
 

tests/integration/test_user_management.py

@@ -1,21 +1,24 @@
 import pytest
 
 from labelbox.schema.organization import ProjectRole
-from conftest import Environ
 
 
 def test_org_invite(client, organization, environ):
     role = client.get_roles()['LABELER']
     dummy_email = "none@labelbox.com"
     invite_limit = organization.invite_limit()
 
-    if environ == Environ.PROD:
+    if environ.value == "prod":
         user_limit = organization.user_limit()
         assert invite_limit.remaining > 0, "No invites available for the account associated with this key."
+    elif environ.value != "staging":
+        raise ValueError(
+            f"Expected tests to run against either prod or staging. Found {environ}"
+        )
 
     invite = organization.invite_user(dummy_email, role)
 
-    if environ == Environ.PROD:
+    if environ.value == "prod":
         invite_limit_after = organization.invite_limit()
         user_limit_after = organization.user_limit()
         # One user added