Changes tests and add refresh cache following review

Author: paulnoirel

libs/labelbox/src/labelbox/client.py

@@ -2466,6 +2466,7 @@ def create_api_key(
         role: Union[Role, str],
         validity: int = 0,
         time_unit: TimeUnit = TimeUnit.SECOND,
+        refresh_cache: bool = False,
     ) -> Dict[str, str]:
         """Creates a new API key.
 
@@ -2475,13 +2476,21 @@ def create_api_key(
             role (Union[Role, str]): The role object or role ID to assign to the API key.
             validity (int, optional): The validity period of the API key. Defaults to 0 (no expiration).
             time_unit (TimeUnit, optional): The time unit for the validity period. Defaults to TimeUnit.SECOND.
+            refresh_cache (bool, optional): Whether to refresh cached permissions and roles. Defaults to False.
 
         Returns:
             Dict[str, str]: A dictionary containing the created API key information.
         """
         warnings.warn(
             "The creation of API keys is currently in alpha and its behavior may change in future releases.",
         )
+        if refresh_cache:
+            # Clear cached attributes if they exist
+            if hasattr(self, "_cached_current_user_permissions"):
+                delattr(self, "_cached_current_user_permissions")
+            if hasattr(self, "_cached_available_api_key_roles"):
+                delattr(self, "_cached_available_api_key_roles")
+
         return ApiKey.create_api_key(
             self, name, user, role, validity, time_unit
         )

libs/labelbox/src/labelbox/schema/api_key.py

@@ -330,24 +330,22 @@ def create_api_key(
             )
 
         validity_seconds = 0
-        if validity > 0:
-            if not isinstance(time_unit, TimeUnit):
-                raise ValueError(
-                    "time_unit must be a valid TimeUnit enum value"
-                )
+        if validity < 0:
+            raise ValueError("validity must be a positive integer")
 
-            validity_seconds = validity * time_unit.value
+        if not isinstance(time_unit, TimeUnit):
+            raise ValueError("time_unit must be a valid TimeUnit enum value")
 
-            if validity_seconds < TimeUnit.MINUTE.value:
-                raise ValueError("Minimum validity period is 1 minute")
+        validity_seconds = validity * time_unit.value
 
-            max_seconds = 25 * TimeUnit.WEEK.value
-            if validity_seconds > max_seconds:
-                raise ValueError(
-                    "Maximum validity period is 6 months (or 25 weeks)"
-                )
-        else:
-            raise ValueError("validity must be a positive integer")
+        if validity_seconds < TimeUnit.MINUTE.value:
+            raise ValueError("Minimum validity period is 1 minute")
+
+        max_seconds = 25 * TimeUnit.WEEK.value
+        if validity_seconds > max_seconds:
+            raise ValueError(
+                "Maximum validity period is 6 months (or 25 weeks)"
+            )
 
         query_str = """
          mutation CreateUserApiKeyPyApi($name: String!, $userEmail: String!, $role: String, $validitySeconds: Int) {
@@ -379,17 +377,7 @@ def create_api_key(
             return api_key_result
 
         except Exception as e:
-            if (
-                "permission" in str(e).lower()
-                or "unauthorized" in str(e).lower()
-            ):
-                raise LabelboxError(
-                    f"Permission denied: You don't have sufficient permissions to create API keys. Original error: {str(e)}"
-                )
-            else:
-                error_message = f"Failed to create API key: {str(e)}"
-                logger.error(error_message)
-                raise LabelboxError(error_message) from e
+            raise LabelboxError(str(e)) from e
 
     @staticmethod
     def get_api_key(client: "Client", api_key_id: str) -> Optional["ApiKey"]:

libs/labelbox/tests/integration/test_api_keys.py

@@ -1,12 +1,56 @@
 import uuid
 import pytest
+import os
 
 from labelbox.schema.timeunit import TimeUnit
 from labelbox.schema.api_key import ApiKey
 from lbox.exceptions import LabelboxError
 # The creation of API keys requires a feature flag to be enabled.
 
 
+@pytest.mark.skipif(
+    condition=os.environ["LABELBOX_TEST_ENVIRON"] != "prod",
+    reason="Admin permissions are required to create API keys",
+)
+def test_create_api_key_success(client):
+    # Create a test API key
+    key_name = f"Test Key {uuid.uuid4()}"
+    user_email = client.get_user().email
+
+    assert (
+        client.get_user().org_role().name == "Admin"
+    ), "User must be an admin to create API keys"
+
+    # Get available roles and use the first one
+    available_roles = ApiKey._get_available_api_key_roles(client)
+    assert (
+        len(available_roles) > 0
+    ), "No available roles found for API key creation"
+
+    # Create the API key with a short validity period
+    api_key_result = client.create_api_key(
+        name=key_name,
+        user=user_email,
+        role=available_roles[0],
+        validity=5,
+        time_unit=TimeUnit.MINUTE,
+    )
+
+    # Verify the response format
+    assert isinstance(
+        api_key_result, dict
+    ), "API key result should be a dictionary"
+    assert "id" in api_key_result, "API key result should contain an 'id' field"
+    assert (
+        "jwt" in api_key_result
+    ), "API key result should contain a 'jwt' field"
+
+    # Verify the JWT token format (should be a JWT string)
+    jwt = api_key_result["jwt"]
+    assert isinstance(jwt, str), "JWT should be a string"
+    assert jwt.count(".") == 2, "JWT should have three parts separated by dots"
+
+
 def test_create_api_key_failed(client):
     # Test with invalid role
     with pytest.raises(ValueError) as excinfo:
@@ -132,7 +176,7 @@ def test_create_api_key_invalid_validity_values(client):
             validity=0,
             time_unit=TimeUnit.MINUTE,
         )
-    assert "validity must be a positive integer" in str(excinfo.value).lower()
+    assert "minimum validity period is 1 minute" in str(excinfo.value).lower()
 
     # Days (exceeding 6 months)
     with pytest.raises(ValueError) as excinfo:
@@ -185,10 +229,16 @@ def test_create_api_key_invalid_time_unit(client):
     assert "valid TimeUnit" in str(excinfo.value)
 
 
+@pytest.mark.skipif(
+    condition=os.environ["LABELBOX_TEST_ENVIRON"] == "prod",
+    reason="Accounts with sdmin permission can create API keys",
+)
 def test_create_api_key_insufficient_permissions(client):
     """Test that creating an API key fails when the user has insufficient permissions."""
     user_email = client.get_user().email
 
+    assert client.get_user().org_role().name == "Admin"
+
     # Attempt to create another API key using the limited permissions client
     # This should fail due to insufficient permissions
     with pytest.raises(LabelboxError) as excinfo:
@@ -200,5 +250,4 @@ def test_create_api_key_insufficient_permissions(client):
             time_unit=TimeUnit.MINUTE,
         )
 
-    # Check for the exact "Permission denied" error message
-    assert "Permission denied" in str(excinfo.value)
+    assert "192" in str(excinfo.value)