Adding node assignment to the task entity (#54)

src/api-service/__app__/onefuzzlib/pools.py

@@ -19,7 +19,7 @@ from onefuzztypes.enums import (
 )
 from onefuzztypes.models import Error
 from onefuzztypes.models import Node as BASE_NODE
-from onefuzztypes.models import NodeCommand
+from onefuzztypes.models import NodeAssignment, NodeCommand
 from onefuzztypes.models import NodeTasks as BASE_NODE_TASK
 from onefuzztypes.models import Pool as BASE_POOL
 from onefuzztypes.models import Scaleset as BASE_SCALESET
@@ -187,6 +187,21 @@ class NodeTasks(BASE_NODE_TASK, ORMMixin):
                 result.append(node)
         return result
 
+    @classmethod
+    def get_node_assignments(cls, task_id: UUID) -> List[NodeAssignment]:
+        result = []
+        for entry in cls.search(query={"task_id": [task_id]}):
+            node = Node.get_by_machine_id(entry.machine_id)
+            if node:
+                node_assignment = NodeAssignment(
+                    node_id=node.machine_id,
+                    scaleset_id=node.scaleset_id,
+                    state=entry.state,
+                )
+                result.append(node_assignment)
+
+        return result
+
     @classmethod
     def get_by_machine_id(cls, machine_id: UUID) -> List["NodeTasks"]:
         return cls.search(query={"machine_id": [machine_id]})

src/api-service/__app__/tasks/__init__.py

@@ -12,6 +12,7 @@ from onefuzztypes.responses import BoolResult
 
 from ..onefuzzlib.heartbeat import Heartbeat
 from ..onefuzzlib.jobs import Job
+from ..onefuzzlib.pools import NodeTasks
 from ..onefuzzlib.request import not_ok, ok, parse_request
 from ..onefuzzlib.task_event import TaskEvent
 from ..onefuzzlib.tasks.config import TaskConfigError, check_config
@@ -71,9 +72,8 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
         task = Task.get_by_task_id(request.task_id)
         if isinstance(task, Error):
             return not_ok(task, context=request.task_id)
-        summary = Heartbeat.get_heartbeats(task.task_id)
+        task.heartbeats = Heartbeat.get_heartbeats(task.task_id)
-        task.heartbeats = summary
+        task.nodes = NodeTasks.get_node_assignments(request.task_id)
-
         task.events = TaskEvent.get_summary(request.task_id)
         return ok(task)
 

src/pytypes/onefuzztypes/models.py

@@ -550,6 +550,12 @@ class TaskEventSummary(BaseModel):
     event_type: str
 
 
+class NodeAssignment(BaseModel):
+    node_id: UUID
+    scaleset_id: Optional[UUID]
+    state: NodeTaskState
+
+
 class Task(BaseModel):
     job_id: UUID
     task_id: UUID = Field(default_factory=uuid4)
@@ -561,3 +567,4 @@ class Task(BaseModel):
     heartbeats: Optional[List[HeartbeatSummary]]
     end_time: Optional[datetime]
     events: Optional[List[TaskEventSummary]]
+    nodes: Optional[List[NodeAssignment]]