Fix duplicate in api_key and imports in test_api_keys

Author: paulnoirel

libs/labelbox/src/labelbox/client.py

@@ -2466,7 +2466,7 @@ def create_api_key(
         role: Union[Role, str],
         validity: int = 0,
         time_unit: TimeUnit = TimeUnit.SECOND,
-    ) -> Dict[str, Any]:
+    ) -> Dict[str, str]:
         """Creates a new API key.
 
         Args:
@@ -2477,7 +2477,7 @@ def create_api_key(
             time_unit (TimeUnit, optional): The time unit for the validity period. Defaults to TimeUnit.SECOND.
 
         Returns:
-            Dict[str, Any]: A dictionary containing the created API key information.
+            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.",

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

@@ -1,5 +1,5 @@
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING
 import logging
 
 from labelbox.orm.db_object import DbObject
@@ -10,6 +10,9 @@
 from labelbox.schema.user import User
 from labelbox.schema.role import Role
 
+if TYPE_CHECKING:
+    from labelbox import Client
+
 logger = logging.getLogger(__name__)
 
 
@@ -115,8 +118,15 @@ def revoke(self) -> Dict[str, Any]:
         return result
 
     @staticmethod
-    def _get_current_user_permissions(client) -> List[str]:
-        """Retrieve current user permissions from the client's organization role with caching."""
+    def _get_current_user_permissions(client: "Client") -> List[str]:
+        """Retrieve current user permissions from the client's organization role with caching.
+
+        Args:
+            client: The Labelbox client instance.
+
+        Returns:
+            List[str]: A list of permission strings associated with the current user's role.
+        """
         if hasattr(client, "_cached_current_user_permissions"):
             return client._cached_current_user_permissions
 
@@ -135,53 +145,9 @@ def _get_current_user_permissions(client) -> List[str]:
         return perms
 
     @staticmethod
-    def get_api_key(client, api_key_id: str) -> Optional["ApiKey"]:
-        """Retrieves a single API key by its ID.
-
-        Args:
-            client: The Labelbox client instance.
-            api_key_id (str): The unique ID of the API key.
-
-        Returns:
-            Optional[ApiKey]: The corresponding ApiKey object if found, otherwise None.
-        """
-        query = """
-        query GetApiKeyPyApi($id: ID!) {
-            apiKey(where: {id: $id}) {
-                id
-                name
-                key
-                createdAt
-                expiresAt
-                revoked
-                user {
-                    id
-                    email
-                }
-                role {
-                    id
-                    name
-                    permissions
-                }
-            }
-        }
-        """
-
-        try:
-            result = client.execute(query, {"id": api_key_id})
-
-            if not result or "apiKey" not in result or not result["apiKey"]:
-                return None
-
-            api_key_data = result["apiKey"]
-            return ApiKey(client, api_key_data)
-
-        except Exception as e:
-            raise LabelboxError(f"Failed to retrieve API key: {str(e)}")
-            return None
-
-    @staticmethod
-    def get_api_keys(client, include_expired: bool = False) -> List["ApiKey"]:
+    def get_api_keys(
+        client: "Client", include_expired: bool = False
+    ) -> List["ApiKey"]:
         """Retrieves all API keys accessible to the current user using the provided client.
 
         Args:
@@ -244,8 +210,22 @@ def get_api_keys(client, include_expired: bool = False) -> List["ApiKey"]:
         return all_keys
 
     @staticmethod
-    def _get_available_api_key_roles(client) -> List[str]:
-        """Get the list of built-in roles available for API key creation with caching."""
+    def _get_available_api_key_roles(client: "Client") -> List[str]:
+        """Get the list of built-in roles available for API key creation with caching.
+
+        This method retrieves all roles available in the organization and filters them
+        based on the current user's permissions. The results are cached on the client
+        to avoid redundant API calls.
+
+        Args:
+            client: The Labelbox client instance.
+
+        Returns:
+            List[str]: A list of role names that can be assigned to API keys.
+
+        Raises:
+            LabelboxError: If there's an error retrieving the user permissions.
+        """
         if hasattr(client, "_cached_available_api_key_roles"):
             return client._cached_available_api_key_roles
         try:
@@ -275,7 +255,7 @@ def _get_available_api_key_roles(client) -> List[str]:
         return available_roles
 
     @staticmethod
-    def _get_user(client, email: str) -> Optional[str]:
+    def _get_user(client: "Client", email: str) -> Optional[str]:
         """Checks if a user with the given email exists in the organization.
 
         Args:
@@ -301,7 +281,7 @@ def _get_user(client, email: str) -> Optional[str]:
 
     @staticmethod
     def create_api_key(
-        client,
+        client: "Client",
         name: str,
         user: Union["User", str],
         role: Union["Role", str],
@@ -319,7 +299,7 @@ def create_api_key(
             time_unit (TimeUnit, optional): Time unit for validity period. Defaults to TimeUnit.SECOND.
 
         Returns:
-            Dict[str, Any]: Dictionary containing the created API key details including id and jwt
+            Dict[str, str]: Dictionary containing the created API key details including id and jwt
 
         Raises:
             ValueError: If invalid parameters are provided
@@ -412,7 +392,7 @@ def create_api_key(
                 raise LabelboxError(error_message) from e
 
     @staticmethod
-    def get_api_key(client, api_key_id: str) -> Optional["ApiKey"]:
+    def get_api_key(client: "Client", api_key_id: str) -> Optional["ApiKey"]:
         """Retrieves a single API key by its ID using the provided client.
 
         Args:

libs/labelbox/tests/integration/test_api_keys.py

@@ -1,12 +1,9 @@
 import uuid
 import pytest
-from datetime import datetime, timezone
 
-from labelbox.schema.role import Role
 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.
 
 
@@ -135,7 +132,7 @@ def test_create_api_key_invalid_validity_values(client):
             validity=0,
             time_unit=TimeUnit.MINUTE,
         )
-    assert "validity" in str(excinfo.value).lower()
+    assert "validity must be a positive integer" in str(excinfo.value).lower()
 
     # Days (exceeding 6 months)
     with pytest.raises(ValueError) as excinfo:
@@ -148,6 +145,17 @@ def test_create_api_key_invalid_validity_values(client):
         )
     assert "maximum validity" in str(excinfo.value).lower()
 
+    # Test with validity period less than 1 minute
+    with pytest.raises(ValueError) as excinfo:
+        client.create_api_key(
+            name=f"Test Key {uuid.uuid4()}",
+            user=user_email,
+            role="Admin",
+            validity=30,  # 30 seconds
+            time_unit=TimeUnit.SECOND,
+        )
+    assert "Minimum validity period is 1 minute" in str(excinfo.value)
+
 
 def test_create_api_key_invalid_time_unit(client):
     """Test that providing invalid time unit causes failure."""