gmicloud 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

gmicloud/_internal/_client/_task_client.py

@@ -1,9 +1,14 @@
+import logging
+from requests.exceptions import RequestException
+
 from ._http_client import HTTPClient
 from ._decorator import handle_refresh_token
 from ._iam_client import IAMClient
 from .._config import TASK_SERVICE_BASE_URL
 from .._models import *
 
+logger = logging.getLogger(__name__)
+
 
 class TaskClient:
     """
@@ -21,17 +26,19 @@ class TaskClient:
         self.iam_client = iam_client
 
     @handle_refresh_token
-    def get_task(self, task_id: str) -> Task:
+    def get_task(self, task_id: str) -> Optional[Task]:
         """
         Retrieves a task from the task service using the given task ID.
 
         :param task_id: The ID of the task to be retrieved.
-        :return: An instance of Task containing the details of the retrieved task.
-        :rtype: Task
+        :return: An instance of Task containing the details of the retrieved task, or None if an error occurs.
         """
-        result = self.client.get("/get_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
-
-        return Task.model_validate(result)
+        try:
+            response = self.client.get("/get_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+            return Task.model_validate(response) if response else None
+        except (RequestException, ValueError) as e:
+            logger.error(f"Failed to retrieve task {task_id}: {e}")
+            return None
 
     @handle_refresh_token
     def get_all_tasks(self) -> GetAllTasksResponse:
@@ -39,70 +46,108 @@ class TaskClient:
         Retrieves all tasks from the task service.
 
         :return: An instance of GetAllTasksResponse containing the retrieved tasks.
-        :rtype: GetAllTasksResponse
         """
-        result = self.client.get("/get_tasks", self.iam_client.get_custom_headers())
-        if not result:
+        try:
+            response = self.client.get("/get_tasks", self.iam_client.get_custom_headers())
+            if not response:
+                logger.error("Empty response from /get_tasks")
+                return GetAllTasksResponse(tasks=[])
+            return GetAllTasksResponse.model_validate(response)
+        except (RequestException, ValueError) as e:
+            logger.error(f"Failed to retrieve all tasks: {e}")
             return GetAllTasksResponse(tasks=[])
 
-        return GetAllTasksResponse.model_validate(result)
-
     @handle_refresh_token
-    def create_task(self, task: Task) -> CreateTaskResponse:
+    def create_task(self, task: Task) -> Optional[CreateTaskResponse]:
         """
         Creates a new task using the provided task object.
 
         :param task: The Task object containing the details of the task to be created.
+        :return: The response object containing created task details, or None if an error occurs.
         """
-        result = self.client.post("/create_task", self.iam_client.get_custom_headers(), task.model_dump())
-
-        return CreateTaskResponse.model_validate(result)
+        try:
+            response = self.client.post("/create_task", self.iam_client.get_custom_headers(), task.model_dump())
+            return CreateTaskResponse.model_validate(response) if response else None
+        except (RequestException, ValueError) as e:
+            logger.error(f"Failed to create task: {e}")
+            return None
 
     @handle_refresh_token
-    def update_task_schedule(self, task: Task):
+    def update_task_schedule(self, task: Task) -> bool:
         """
         Updates the schedule of an existing task.
 
         :param task: The Task object containing the updated task details.
+        :return: True if update is successful, False otherwise.
         """
-        self.client.put("/update_schedule", self.iam_client.get_custom_headers(), task.model_dump())
+        try:
+            response = self.client.put("/update_schedule", self.iam_client.get_custom_headers(), task.model_dump())
+            return response is not None
+        except RequestException as e:
+            logger.error(f"Failed to update schedule for task {task.task_id}: {e}")
+            return False
 
     @handle_refresh_token
-    def start_task(self, task_id: str):
+    def start_task(self, task_id: str) -> bool:
         """
         Starts a task using the given task ID.
 
         :param task_id: The ID of the task to be started.
+        :return: True if start is successful, False otherwise.
         """
-        self.client.post("/start_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+        try:
+            response = self.client.post("/start_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+            return response is not None
+        except RequestException as e:
+            logger.error(f"Failed to start task {task_id}: {e}")
+            return False
 
     @handle_refresh_token
-    def stop_task(self, task_id: str):
+    def stop_task(self, task_id: str) -> bool:
         """
         Stops a running task using the given task ID.
 
         :param task_id: The ID of the task to be stopped.
+        :return: True if stop is successful, False otherwise.
         """
-        self.client.post("/stop_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+        try:
+            response = self.client.post("/stop_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+            return response is not None
+        except RequestException as e:
+            logger.error(f"Failed to stop task {task_id}: {e}")
+            return False
 
     @handle_refresh_token
-    def get_usage_data(self, start_timestamp: str, end_timestamp: str) -> GetUsageDataResponse:
+    def get_usage_data(self, start_timestamp: str, end_timestamp: str) -> Optional[GetUsageDataResponse]:
         """
         Retrieves the usage data of a task using the given task ID.
 
         :param start_timestamp: The start timestamp of the usage data.
         :param end_timestamp: The end timestamp of the usage data.
-        """
-        result = self.client.get("/get_usage_data", self.iam_client.get_custom_headers(),
-                                 {"start_timestamp": start_timestamp, "end_timestamp": end_timestamp})
-
-        return result
+        :return: An instance of GetUsageDataResponse, or None if an error occurs.
+        """
+        try:
+            response = self.client.get(
+                "/get_usage_data",
+                self.iam_client.get_custom_headers(),
+                {"start_timestamp": start_timestamp, "end_timestamp": end_timestamp}
+            )
+            return GetUsageDataResponse.model_validate(response) if response else None
+        except (RequestException, ValueError) as e:
+            logger.error(f"Failed to retrieve usage data from {start_timestamp} to {end_timestamp}: {e}")
+            return None
 
     @handle_refresh_token
-    def archive_task(self, task_id: str):
+    def archive_task(self, task_id: str) -> bool:
         """
         Archives a task using the given task ID.
 
         :param task_id: The ID of the task to be archived.
-        """
-        self.client.post("/archive_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+        :return: True if archiving is successful, False otherwise.
+        """
+        try:
+            response = self.client.post("/archive_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+            return response is not None
+        except RequestException as e:
+            logger.error(f"Failed to archive task {task_id}: {e}")
+            return False

gmicloud/_internal/_enums.py

@@ -23,3 +23,11 @@ class TaskEndpointStatus(str, Enum):
     READY = "ready"
     UNREADY = "unready"
     NEW = "new"
+
+class TaskStatus(str, Enum):
+    IDLE = "idle"
+    STARTING = "starting"
+    IN_QUEUE = "in-queue"
+    RUNNING = "running"
+    NEEDSTOP = "needstop"
+    ARCHIVED = "archived"

gmicloud/_internal/_manager/_artifact_manager.py

@@ -81,7 +81,11 @@ class ArtifactManager:
         if not artifact_template_id or not artifact_template_id.strip():
             raise ValueError("Artifact template ID is required and cannot be empty.")
 
-        return self.artifact_client.create_artifact_from_template(artifact_template_id).artifact_id
+        resp = self.artifact_client.create_artifact_from_template(artifact_template_id)
+        if not resp or not resp.artifact_id:
+            raise ValueError("Failed to create artifact from template.")
+
+        return resp.artifact_id
 
     def rebuild_artifact(self, artifact_id: str) -> RebuildArtifactResponse:
         """
@@ -170,7 +174,11 @@ class ArtifactManager:
 
         req = GetBigFileUploadUrlRequest(artifact_id=artifact_id, file_name=model_file_name, file_type=model_file_type)
 
-        return self.artifact_client.get_bigfile_upload_url(req).upload_link
+        resp = self.artifact_client.get_bigfile_upload_url(req)
+        if not resp or not resp.upload_link:
+            raise ValueError("Failed to get bigfile upload URL.")
+
+        return resp.upload_link
 
     def delete_bigfile(self, artifact_id: str, file_name: str) -> str:
         """
@@ -182,7 +190,11 @@ class ArtifactManager:
         self._validate_artifact_id(artifact_id)
         self._validate_file_name(file_name)
 
-        return self.artifact_client.delete_bigfile(artifact_id, file_name).status
+        resp = self.artifact_client.delete_bigfile(artifact_id, file_name)
+        if not resp or not resp.status:
+            raise ValueError("Failed to delete bigfile.")
+
+        return resp.status
 
     def upload_large_file(self, upload_link: str, file_path: str) -> None:
         """
@@ -229,14 +241,14 @@ class ArtifactManager:
 
         return artifact_id
 
-    def get_artifact_templates(self) -> List[ArtifactTemplate]:
+    def get_public_templates(self) -> List[ArtifactTemplate]:
         """
         Fetch all artifact templates.
 
         :return: A list of ArtifactTemplate objects.
         :rtype: List[ArtifactTemplate]
         """
-        return self.artifact_client.get_artifact_templates().artifact_templates
+        return self.artifact_client.get_public_templates()
 
     @staticmethod
     def _validate_file_name(file_name: str) -> None:

gmicloud/_internal/_manager/_iam_manager.py

@@ -0,0 +1,36 @@
+from datetime import datetime
+from .._client._iam_client import IAMClient
+from .._models import *
+
+
+class IAMManager:
+    """
+    IamManager handles operations related to IAM, including user authentication and authorization.
+    """
+
+    def __init__(self, iam_client: IAMClient):
+        """
+        Initialize the IAMManager instance and the associated IAMClient.
+        """
+        self.iam_client = iam_client
+
+    def create_org_api_key(self, name: str, expires_at: Optional[int] = None) -> str:
+        """
+        Creates a new API key for the current user.
+        """
+
+        if not name:
+            raise ValueError("API key name cannot be empty")
+        if not expires_at:
+            # Set the expiration date to 30 days from now
+            expires_at = int(datetime.now().timestamp()) + 30 * 24 * 60 * 60
+            print(expires_at)
+
+        return self.iam_client.create_org_api_key(
+            CreateAPIKeyRequest(name=name, type="ie_model", expiresAt=expires_at))
+
+    def get_org_api_keys(self) -> List[APIKey]:
+        """
+        Fetches all API keys for the current user.
+        """
+        return self.iam_client.get_org_api_keys().keys

gmicloud/_internal/_manager/_task_manager.py

@@ -37,7 +37,11 @@ class TaskManager:
 
         :return: A list of `Task` objects.
         """
-        return self.task_client.get_all_tasks(self.iam_client.get_user_id()).tasks
+        resp = self.task_client.get_all_tasks(self.iam_client.get_user_id())
+        if not resp or not resp.tasks:
+            return []
+
+        return resp.tasks
 
     def create_task(self, task: Task) -> Task:
         """
@@ -51,8 +55,11 @@ class TaskManager:
         self._validate_task(task)
         if not task.owner:
             task.owner = TaskOwner(user_id=self.iam_client.get_user_id())
+        resp = self.task_client.create_task(task)
+        if not resp or not resp.task:
+            raise ValueError("Failed to create task.")
 
-        return self.task_client.create_task(task).task
+        return resp.task
 
     def create_task_from_file(self, artifact_id: str, config_file_path: str, trigger_timestamp: int = None) -> Task:
         """
@@ -76,7 +83,7 @@ class TaskManager:
 
         return self.create_task(task)
 
-    def update_task_schedule(self, task: Task):
+    def update_task_schedule(self, task: Task) -> bool:
         """
         Update the schedule of an existing task.
 
@@ -87,10 +94,10 @@ class TaskManager:
         self._validate_task(task)
         self._validate_not_empty(task.task_id, "Task ID")
 
-        self.task_client.update_task_schedule(task)
+        return self.task_client.update_task_schedule(task)
 
     def update_task_schedule_from_file(self, artifact_id: str, task_id: str, config_file_path: str,
-                                       trigger_timestamp: int = None):
+                                       trigger_timestamp: int = None) -> bool:
         """
         Update the schedule of an existing task using data from a file. The file should contain a valid task definition.
 
@@ -112,9 +119,9 @@ class TaskManager:
         if trigger_timestamp:
             task.config.task_scheduling.scheduling_oneoff.trigger_timestamp = trigger_timestamp
 
-        self.update_task_schedule(task)
+        return self.update_task_schedule(task)
 
-    def start_task(self, task_id: str):
+    def start_task(self, task_id: str) -> bool:
         """
         Start a task by its ID.
 
@@ -124,9 +131,9 @@ class TaskManager:
         """
         self._validate_not_empty(task_id, "Task ID")
 
-        self.task_client.start_task(task_id)
+        return self.task_client.start_task(task_id)
 
-    def stop_task(self, task_id: str):
+    def stop_task(self, task_id: str) -> bool:
         """
         Stop a task by its ID.
 
@@ -136,7 +143,7 @@ class TaskManager:
         """
         self._validate_not_empty(task_id, "Task ID")
 
-        self.task_client.stop_task(task_id)
+        return self.task_client.stop_task(task_id)
 
     def get_usage_data(self, start_timestamp: str, end_timestamp: str) -> GetUsageDataResponse:
         """
@@ -151,7 +158,7 @@ class TaskManager:
 
         return self.task_client.get_usage_data(start_timestamp, end_timestamp)
 
-    def archive_task(self, task_id: str):
+    def archive_task(self, task_id: str) -> bool:
         """
         Archive a task by its ID.
 
@@ -161,7 +168,7 @@ class TaskManager:
         """
         self._validate_not_empty(task_id, "Task ID")
 
-        self.task_client.archive_task(task_id)
+        return self.task_client.archive_task(task_id)
 
     @staticmethod
     def _validate_not_empty(value: str, name: str):

gmicloud/_internal/_models.py

@@ -2,7 +2,7 @@ from typing import Optional, List
 from datetime import datetime
 
 from pydantic import BaseModel
-from gmicloud._internal._enums import BuildStatus, TaskEndpointStatus
+from gmicloud._internal._enums import BuildStatus, TaskStatus, TaskEndpointStatus
 
 
 class BigFileMetadata(BaseModel):
@@ -24,6 +24,7 @@ class ArtifactMetadata(BaseModel):
     artifact_description: Optional[str] = ""  # Description of the artifact.
     artifact_tags: Optional[List[str]] = ""  # Comma-separated tags for categorizing the artifact.
     artifact_volume_path: Optional[str] = ""  # Path to the volume where the artifact is stored.
+    artifact_template_id: Optional[str] = ""  # The template ID used to create this artifact.
 
 
 class ArtifactData(BaseModel):
@@ -130,7 +131,7 @@ class DeleteBigfileResponse(BaseModel):
     status: Optional[str] = ""  # Status of the deletion process.
 
 
-class GetArtifactTemplatesResponse(BaseModel):
+class GetPublicTemplatesResponse(BaseModel):
     """
     Response containing a list of artifact templates.
     """
@@ -141,24 +142,46 @@ class ArtifactTemplate(BaseModel):
     """
     Template for creating an artifact.
     """
-    artifact_template_id: str  # Unique identifier for the artifact template.
-    artifact_description: Optional[str] = ""  # Description of the artifact template.
-    artifact_name: Optional[str] = ""  # Name of the artifact template.
-    artifact_tags: Optional[List[str]] = None  # Tags associated with the artifact template.
-    ray: Optional["RayTemplate"] = None  # Template for Ray-based artifacts.
+    template_id: str  # Unique identifier for the artifact template.
+    template_data: Optional["TemplateData"] = None  # Data for the artifact template.
+    template_metadata: Optional["TemplateMetadata"] = None  # Metadata for the artifact template.
+
+
+class TemplateMetadata(BaseModel):
+    """
+    Metadata for an artifact template.
+    """
+    create_at: Optional[str] = None  # Timestamp when the template was created.
+    create_by: Optional[str] = ""  # ID of the user who created the template.
+    create_by_org_id: Optional[str] = ""  # ID of the organization to which the user belongs.
+    is_public: Optional[bool] = False  # Indicates if the template is public.
+    update_at: Optional[str] = None  # Timestamp when the template was last updated.
+    update_by: Optional[str] = ""  # ID of the user who last updated the template.
+
+class TemplateData(BaseModel):
+    """
+    Data for an artifact template.
+    """
+    description: Optional[str] = ""  # Description of the artifact template.
+    icon_link: Optional[str] = ""  # Link to the icon for the artifact template.
+    image_link: Optional[str] = ""  # Link to the image for the artifact template.
+    name: Optional[str] = ""  # Name of the artifact template.
+    ray: Optional["RayContent"] = None  # Template for Ray-based artifacts.
     resources: Optional["ResourcesTemplate"] = None  # Resource allocation template.
+    tags: Optional[List[str]] = None  # Tags associated with the artifact template.
+    volume_path: Optional[str] = ""  # Path to the volume where the artifact is stored.
 
 
-class RayTemplate(BaseModel):
+class RayContent(BaseModel):
     deployment_name: Optional[str] = ""  # Name of the deployment.
     file_path: Optional[str] = ""  # Path to the task file in storage.
-    version: Optional[str] = ""  # Version of Ray used.
 
 
 class ResourcesTemplate(BaseModel):
     cpu: Optional[int] = 0  # Number of CPU cores allocated.
     memory: Optional[int] = 0  # Amount of RAM (in GB) allocated.
     gpu: Optional[int] = 0  # Number of GPUs allocated.
+    gpu_name: Optional[str] = ""  # Type the GPU allocated.
 
 
 class CreateArtifactFromTemplateRequest(BaseModel):
@@ -290,7 +313,7 @@ class Task(BaseModel):
     config: Optional[TaskConfig] = None  # Configuration data for the task.
     endpoint_info: Optional[EndpointInfo] = None  # Additional information about the task endpoint.
     cluster_endpoints: Optional[List[EndpointInfo]] = None  # Endpoints for the task cluster.
-    task_status: Optional[str] = ""  # Status of the task.
+    task_status: Optional[TaskStatus] = ""  # Status of the task.
     readiness_status: Optional[str] = ""  # Readiness status of the task.
     user_preference: Optional[UserPreference] = None  # User preference for the task.
 
@@ -373,3 +396,73 @@ class LoginRequest(BaseModel):
     """
     email: str  # User email.
     password: str  # User password.
+
+
+class User(BaseModel):
+    """
+    User information.
+    """
+    id: Optional[str] = ""  # User ID.
+    email: Optional[str] = ""  # User email.
+    firstName: Optional[str] = ""  # User first name.
+    lastName: Optional[str] = ""  # User last name.
+
+
+class Organization(BaseModel):
+    """
+    Organization information.
+    """
+    id: Optional[str] = ""  # Organization ID.
+    role: Optional[str] = ""  # Organization role.
+
+
+class ProfileResponse(BaseModel):
+    """
+    Response object for user profile.
+    """
+    user: User  # User information.
+    organization: Organization  # Organization information.
+
+
+class CreateAPIKeyRequest(BaseModel):
+    """
+    Request object for creating an API key.
+    """
+    name: str  # Name of the API key.
+    type: Optional[str] = ""  # Type of the API key.
+    expiresAt: Optional[int] = 0  # Expiration timestamp for the API key.
+
+
+class CreateAPIKeyResponse(BaseModel):
+    """
+    Response object for creating an API key.
+    """
+    key: str  # The created API key.
+
+
+class APIKey(BaseModel):
+    """
+    API key information.
+    """
+    id: Optional[str] = ""  # API key ID.
+    name: Optional[str] = ""  # API key name.
+    type: Optional[str] = ""  # API key type.
+    partialKey: Optional[str] = ""  # Partial key for the API key.
+    expiresAt: Optional[int] = 0  # Expiration timestamp for the API key.
+    createdAt: Optional[int] = 0  # Creation timestamp for the API key.
+    owner: Optional[User] = None  # Owner of the API key.
+
+
+class GetAPIKeysResponse(BaseModel):
+    """
+    Response object for getting a list of API keys.
+    """
+    keys: list[APIKey]  # List of API keys.
+
+
+class GetSelfAPIKeyResponse(BaseModel):
+    """
+    Response object for getting the API key of the current user.
+    """
+    key: APIKey  # The API key of the current user.
+    organization: Optional[Organization] = None  # Organization information.

gmicloud/client.py

@@ -1,14 +1,18 @@
 import os
 import time
+import logging
 
 from typing import Optional
 
 from ._internal._client._iam_client import IAMClient
 from ._internal._manager._artifact_manager import ArtifactManager
 from ._internal._manager._task_manager import TaskManager
+from ._internal._manager._iam_manager import IAMManager
 from ._internal._enums import BuildStatus
 from ._internal._models import Task, TaskConfig, RayTaskConfig, TaskScheduling, ReplicaResource
 
+logger = logging.getLogger(__name__)
+
 
 class Client:
     def __init__(self, client_id: Optional[str] = "", email: Optional[str] = "", password: Optional[str] = ""):
@@ -32,6 +36,7 @@ class Client:
         # Managers are lazily initialized through private attributes
         self._artifact_manager = None
         self._task_manager = None
+        self._iam_manager = None
 
     def create_task_from_artifact_template(self, artifact_template_id: str, task_scheduling: TaskScheduling) -> Task:
         """
@@ -50,24 +55,28 @@ class Client:
         artifact_manager = self.artifact_manager
         task_manager = self.task_manager
 
-        templates = artifact_manager.get_artifact_templates()
+        templates = artifact_manager.get_public_templates()
         template = None
         for v in templates:
-            if v.artifact_template_id == artifact_template_id:
+            if v.template_id == artifact_template_id:
                 template = v
         if not template:
             raise ValueError(f"Template with ID {artifact_template_id} not found.")
-        if not template.ray:
+        if not template.template_data:
+            raise ValueError("Template does not contain template data.")
+        if not template.template_data.ray:
             raise ValueError("Template does not contain Ray configuration.")
-        if not template.resources:
+        if not template.template_data.resources:
             raise ValueError("Template does not contain resource configuration.")
 
         artifact_id = artifact_manager.create_artifact_from_template(artifact_template_id)
+
+        logger.info(f"Successfully created artifact from template, artifact_id: {artifact_id}")
         # Wait for the artifact to be ready
         while True:
             try:
                 artifact = artifact_manager.get_artifact(artifact_id)
-                print(f"Artifact status: {artifact.build_status}")
+                logger.info(f"Successfully got artifact info, artifact status: {artifact.build_status}")
                 # Wait until the artifact is ready
                 if artifact.build_status == BuildStatus.SUCCESS:
                     break
@@ -94,8 +103,10 @@ class Client:
                 ),
             ))
 
+            logger.info(f"Successfully created task, task_id: {task.task_id}")
             # Start the task
             task_manager.start_task(task.task_id)
+            logger.info(f"Successfully started task, task_id: {task.task_id}")
         except Exception as e:
             raise e
 
@@ -120,3 +131,13 @@ class Client:
         if self._task_manager is None:
             self._task_manager = TaskManager(self.iam_client)
         return self._task_manager
+
+    @property
+    def iam_manager(self):
+        """
+        Lazy initialization for IAMManager.
+        Ensures the Client instance controls its lifecycle.
+        """
+        if self._iam_manager is None:
+            self._iam_manager = IAMManager(self.iam_client)
+        return self._iam_manager

gmicloud/tests/test_artifacts.py

@@ -251,24 +251,23 @@ class TestArtifactManager(unittest.TestCase):
             self.artifact_manager.delete_bigfile("nonexistent_id", "file.txt")
         self.assertTrue("Artifact not found" in str(context.exception))
 
-    @patch('gmicloud._internal._client._artifact_client.ArtifactClient.get_artifact_templates')
-    def test_get_artifact_templates_returns_templates(self, mock_get_artifact_templates):
-        mock_get_artifact_templates.return_value = GetArtifactTemplatesResponse(
-            artifact_templates=[ArtifactTemplate(artifact_template_id="1", artifact_name="Template1")])
-        templates = self.artifact_manager.get_artifact_templates()
+    @patch('gmicloud._internal._client._artifact_client.ArtifactClient.get_public_templates')
+    def test_get_artifact_templates_returns_templates(self, mock_get_public_templates):
+        mock_get_public_templates.return_value = [ArtifactTemplate(template_id="1", template_data=TemplateData(name="Template1"))]
+        templates = self.artifact_manager.get_public_templates()
         self.assertEqual(len(templates), 1)
-        self.assertEqual(templates[0].artifact_template_id, "1")
-        self.assertEqual(templates[0].artifact_name, "Template1")
+        self.assertEqual(templates[0].template_id, "1")
+        self.assertEqual(templates[0].template_data.name, "Template1")
 
-    @patch('gmicloud._internal._client._artifact_client.ArtifactClient.get_artifact_templates')
-    def test_get_artifact_templates_returns_empty_list_when_no_templates(self, mock_get_artifact_templates):
-        mock_get_artifact_templates.return_value = GetArtifactTemplatesResponse(artifact_templates=[])
-        templates = self.artifact_manager.get_artifact_templates()
+    @patch('gmicloud._internal._client._artifact_client.ArtifactClient.get_public_templates')
+    def test_get_artifact_templates_returns_empty_list_when_no_templates(self, mock_get_public_templates):
+        mock_get_public_templates.return_value = []
+        templates = self.artifact_manager.get_public_templates()
         self.assertEqual(len(templates), 0)
 
-    @patch('gmicloud._internal._client._artifact_client.ArtifactClient.get_artifact_templates')
-    def test_get_artifact_templates_raises_error_on_failure(self, mock_get_artifact_templates):
-        mock_get_artifact_templates.side_effect = Exception("Failed to fetch templates")
+    @patch('gmicloud._internal._client._artifact_client.ArtifactClient.get_public_templates')
+    def test_get_artifact_templates_raises_error_on_failure(self, mock_get_public_templates):
+        mock_get_public_templates.side_effect = Exception("Failed to fetch templates")
         with self.assertRaises(Exception) as context:
-            self.artifact_manager.get_artifact_templates()
+            self.artifact_manager.get_public_templates()
         self.assertTrue("Failed to fetch templates" in str(context.exception))

gmicloud/tests/test_tasks.py

@@ -72,7 +72,7 @@ class TestTaskManager(unittest.TestCase):
                 'create_timestamp': 1734951094,
                 'last_update_timestamp': 1734975798
             },
-            'task_status': '',
+            'task_status': 'running',
             'readiness_status': '',
             'info': {
                 'endpoint_status': 'ready',