refactor: update model names for clarification

PrivateAIM · Mar 11, 2024 · a8a5c7e · a8a5c7e
1 parent 1670fc0
commit a8a5c7e
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 17 deletions.
diff --git a/gateway/models.py b/gateway/models.py
@@ -182,6 +182,17 @@ class ApprovalStatus(Enum):
     rejected: str = "rejected"
 
 
+class AnalysisNodeRunStatus(Enum):
+    """Possible values for analysis run status."""
+    starting: str = "starting"
+    started: str = "started"
+    running: str = "running"
+    stopping: str = "stopping"
+    stopped: str = "stopped"
+    finished: str = "finished"
+    failed: str = "failed"
+
+
 ## Response Models
 class BaseHubResponse(BaseModel):
     """Common attributes of Hub responses."""
@@ -200,7 +211,7 @@ class MasterImage(BaseHubResponse):
     command_arguments: str | None = None
 
 
-class ProjectResponse(BaseHubResponse):
+class Project(BaseHubResponse):
     """Single project response model."""
     name: str
     analyses: int
@@ -212,10 +223,10 @@ class ProjectResponse(BaseHubResponse):
 
 class AllProjects(BaseModel):
     """List of all projects."""
-    data: list[ProjectResponse]
+    data: list[Project]
 
 
-class NodeDetails(BaseHubResponse):
+class Node(BaseHubResponse):
     """Node details."""
     external_name: str | None = None
     name: str
@@ -228,7 +239,7 @@ class NodeDetails(BaseHubResponse):
     realm_id: uuid.UUID
 
 
-class AnalysisOrProjectNodeResponse(BaseHubResponse):
+class AnalysisOrProjectNode(BaseHubResponse):
     """Single project or analysis by node."""
 
     approval_status: ApprovalStatus
@@ -239,20 +250,20 @@ class AnalysisOrProjectNodeResponse(BaseHubResponse):
     node_realm_id: uuid.UUID | None = None
 
 
-class ListAnalysisOrProjectNodeResponse(BaseModel):
-    data: list[AnalysisOrProjectNodeResponse]
+class ListAnalysisOrProjectNodes(BaseModel):
+    data: list[AnalysisOrProjectNode]
 
 
-class AnalysisNodeResponse(AnalysisOrProjectNodeResponse):
+class AnalysisNode(AnalysisOrProjectNode):
     """Node analysis response model."""
-    run_status: str | None = None
+    run_status: AnalysisNodeRunStatus
     index: int
     artifact_tag: str | None = None
     artifact_digest: str | None = None
     analysis_id: uuid.UUID
     analysis_realm_id: uuid.UUID
-    node: NodeDetails | None = None
+    node: Node | None = None
 
 
-class ListAnalysisNodeResponse(BaseModel):
-    data: list[AnalysisNodeResponse]
+class ListAnalysisNodes(BaseModel):
+    data: list[AnalysisNode]
diff --git a/gateway/routers/hub.py b/gateway/routers/hub.py
@@ -10,8 +10,8 @@
 from gateway.auth import hub_oauth2_scheme
 from gateway.conf import gateway_settings
 from gateway.core import route
-from gateway.models import ImageDataResponse, ContainerResponse, ProjectResponse, AllProjects, \
-    ApprovalStatus, AnalysisOrProjectNodeResponse, ListAnalysisNodeResponse, ListAnalysisOrProjectNodeResponse
+from gateway.models import ImageDataResponse, ContainerResponse, Project, AllProjects, \
+    ApprovalStatus, AnalysisOrProjectNode, ListAnalysisNodes, ListAnalysisOrProjectNodes
 
 hub_router = APIRouter(
     dependencies=[Security(hub_oauth2_scheme)],
@@ -118,7 +118,7 @@ async def list_all_projects(
     path="/projects/{project_id}",
     status_code=status.HTTP_200_OK,
     service_url=gateway_settings.HUB_SERVICE_URL,
-    response_model=ProjectResponse,
+    response_model=Project,
 )
 async def list_specific_project(
         project_id: Annotated[uuid.UUID, Path(description="Project UUID.")],
@@ -134,7 +134,7 @@ async def list_specific_project(
     path="/project-nodes",
     status_code=status.HTTP_200_OK,
     service_url=gateway_settings.HUB_SERVICE_URL,
-    response_model=ListAnalysisOrProjectNodeResponse,
+    response_model=ListAnalysisOrProjectNodes,
     query_params=["filter_id", "filter_approval_status", "filter_project_id", "filter_project_realm_id",
                   "filter_node_id", "filter_node_realm_id"],
 )
@@ -187,7 +187,7 @@ async def list_project_node(
     path="/project-nodes",
     status_code=status.HTTP_200_OK,
     service_url=gateway_settings.HUB_SERVICE_URL,
-    response_model=AnalysisOrProjectNodeResponse,
+    response_model=AnalysisOrProjectNode,
     body_params=["project_id", "node_id"],
     query_params=["approval_status"],
 )
@@ -209,7 +209,7 @@ async def create_project_node(
     path="/analysis-nodes",
     status_code=status.HTTP_200_OK,
     service_url=gateway_settings.HUB_SERVICE_URL,
-    response_model=ListAnalysisNodeResponse,
+    response_model=ListAnalysisNodes,
     query_params=["filter_id", "filter_approval_status", "filter_project_id", "filter_project_realm_id",
                   "filter_node_id", "filter_node_realm_id", "include"],
 )