mirror of
https://github.com/microsoft/onefuzz.git
synced 2025-06-16 03:48:09 +00:00
Adding node assignment to the task entity (#54)
This commit is contained in:
@ -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]})
|
||||
|
Reference in New Issue
Block a user