Adding node assignment to the task entity (#54)

This commit is contained in:
Cheick Keita
2020-09-29 16:54:17 -07:00
committed by GitHub
parent 6cef9b1bee
commit 932cf7c44d
3 changed files with 26 additions and 4 deletions

View File

@ -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]})