strepsiades
114 lines
3.5 KiB
Python
114 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, List, Optional
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from sync_state_machine.common.sync_node import SyncNode
|
|
from sync_state_machine.datasource.datasource import BaseDataSource
|
|
from sync_state_machine.datasource.task_result import TaskResult
|
|
|
|
|
|
class MockSchema(BaseModel):
|
|
id: Optional[str] = None
|
|
name: str = "mock"
|
|
|
|
|
|
class MockSyncNode(SyncNode[MockSchema]):
|
|
node_type = "mock"
|
|
schema = MockSchema
|
|
|
|
|
|
@dataclass
|
|
class HandlerScript:
|
|
create: str = "success" # success | failed | skipped | in_progress
|
|
update: str = "success"
|
|
delete: str = "success"
|
|
poll_sequence: List[str] = field(default_factory=lambda: ["success"]) # success | failed | in_progress
|
|
|
|
|
|
class ScriptedMockHandler:
|
|
def __init__(self, script: HandlerScript):
|
|
self._script = script
|
|
self._collection = None
|
|
self._poll_index = 0
|
|
|
|
@property
|
|
def node_type(self) -> str:
|
|
return "mock"
|
|
|
|
@property
|
|
def node_class(self):
|
|
return MockSyncNode
|
|
|
|
@property
|
|
def schema(self):
|
|
return MockSchema
|
|
|
|
def set_collection(self, collection):
|
|
self._collection = collection
|
|
|
|
async def load(self):
|
|
return []
|
|
|
|
async def create_all(self, nodes):
|
|
return [self._build_result(self._script.create, node.node_id, "create") for node in nodes]
|
|
|
|
async def update_all(self, nodes):
|
|
return [
|
|
self._build_result(
|
|
self._script.update,
|
|
node.node_id,
|
|
"update",
|
|
)
|
|
for node in nodes
|
|
]
|
|
|
|
async def delete_all(self, nodes):
|
|
return [self._build_result(self._script.delete, node.node_id, "delete") for node in nodes]
|
|
|
|
async def poll_tasks(self, task_ids: List[str]) -> Dict[str, TaskResult]:
|
|
status = self._script.poll_sequence[min(self._poll_index, len(self._script.poll_sequence) - 1)]
|
|
self._poll_index += 1
|
|
out: Dict[str, TaskResult] = {}
|
|
for task_id in task_ids:
|
|
out[task_id] = self._build_poll_result(status, task_id)
|
|
return out
|
|
|
|
def _build_result(
|
|
self,
|
|
mode: str,
|
|
node_id: str,
|
|
phase: str,
|
|
) -> TaskResult:
|
|
if mode == "success":
|
|
return TaskResult.success(node_id=node_id, data_id=f"{phase}-{node_id}")
|
|
if mode == "failed":
|
|
return TaskResult.failed(node_id=node_id, error=f"{phase} failed")
|
|
if mode == "skipped":
|
|
return TaskResult.skipped(
|
|
node_id=node_id,
|
|
reason=f"{phase} skipped",
|
|
)
|
|
if mode == "in_progress":
|
|
return TaskResult.in_progress(node_id=node_id, task_id=f"task-{phase}-{node_id}")
|
|
raise ValueError(f"Unsupported mode: {mode}")
|
|
|
|
def _build_poll_result(self, mode: str, task_id: str) -> TaskResult:
|
|
if mode == "success":
|
|
return TaskResult.success(node_id=self._node_id_from_task(task_id), data_id=f"done-{task_id}")
|
|
if mode == "failed":
|
|
return TaskResult.failed(node_id=self._node_id_from_task(task_id), error="poll failed")
|
|
if mode == "in_progress":
|
|
return TaskResult.in_progress(node_id=self._node_id_from_task(task_id), task_id=task_id)
|
|
raise ValueError(f"Unsupported poll mode: {mode}")
|
|
|
|
@staticmethod
|
|
def _node_id_from_task(task_id: str) -> str:
|
|
return task_id.split("-")[-1]
|
|
|
|
|
|
class MockDataSource(BaseDataSource):
|
|
pass
|