1use anyhow::Result;
11use async_trait::async_trait;
12
13#[async_trait]
18pub trait WorkerHandle: Send + 'static {
19 fn id(&self) -> u32;
21
22 async fn ready(&mut self) -> Result<()>;
25
26 async fn execute(
28 &mut self,
29 method: &str,
30 payload: serde_json::Value,
31 ) -> Result<serde_json::Value>;
32
33 async fn terminate(&mut self) -> Result<()>;
36}
37
38#[async_trait]
40pub trait Runtime: Send + Sync + 'static {
41 async fn spawn(&self) -> Result<Box<dyn WorkerHandle>>;
45}
46
47type MockResponder =
50 std::sync::Arc<dyn Fn(&str, &serde_json::Value) -> Result<serde_json::Value> + Send + Sync>;
51
52pub struct MockRuntime {
54 responder: MockResponder,
55 next_id: std::sync::atomic::AtomicU32,
56}
57
58impl MockRuntime {
59 pub fn echo() -> Self {
61 Self {
62 responder: std::sync::Arc::new(|_method, payload| Ok(payload.clone())),
63 next_id: std::sync::atomic::AtomicU32::new(10000),
64 }
65 }
66}
67
68#[async_trait]
69impl Runtime for MockRuntime {
70 async fn spawn(&self) -> Result<Box<dyn WorkerHandle>> {
71 let id = self
72 .next_id
73 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
74 Ok(Box::new(MockWorker {
75 id,
76 responder: self.responder.clone(),
77 terminated: false,
78 }))
79 }
80}
81
82pub struct MockWorker {
84 id: u32,
85 responder: MockResponder,
86 terminated: bool,
87}
88
89#[async_trait]
90impl WorkerHandle for MockWorker {
91 fn id(&self) -> u32 {
92 self.id
93 }
94
95 async fn ready(&mut self) -> Result<()> {
96 Ok(())
97 }
98
99 async fn execute(
100 &mut self,
101 method: &str,
102 payload: serde_json::Value,
103 ) -> Result<serde_json::Value> {
104 if self.terminated {
105 anyhow::bail!("worker terminated");
106 }
107 (self.responder)(method, &payload)
108 }
109
110 async fn terminate(&mut self) -> Result<()> {
111 self.terminated = true;
112 Ok(())
113 }
114}