common_procedure_test/
lib.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Test utilities for procedures.
16
17use std::collections::HashMap;
18use std::sync::Arc;
19
20use async_trait::async_trait;
21use common_procedure::local::{DynamicKeyLockGuard, acquire_dynamic_key_lock};
22use common_procedure::rwlock::KeyRwLock;
23use common_procedure::store::poison_store::PoisonStore;
24use common_procedure::test_util::InMemoryPoisonStore;
25use common_procedure::{
26    Context, ContextProvider, Output, PoisonKey, Procedure, ProcedureId, ProcedureState,
27    ProcedureWithId, Result, Status, StringKey,
28};
29use tokio::sync::watch::Receiver;
30
31/// A Mock [ContextProvider].
32#[derive(Default)]
33pub struct MockContextProvider {
34    states: HashMap<ProcedureId, ProcedureState>,
35    poison_manager: InMemoryPoisonStore,
36    dynamic_key_lock: Arc<KeyRwLock<String>>,
37}
38
39impl MockContextProvider {
40    /// Returns a new provider.
41    pub fn new(states: HashMap<ProcedureId, ProcedureState>) -> MockContextProvider {
42        MockContextProvider {
43            states,
44            poison_manager: InMemoryPoisonStore::default(),
45            dynamic_key_lock: Arc::new(KeyRwLock::new()),
46        }
47    }
48
49    /// Returns a reference to the poison manager.
50    pub fn poison_manager(&self) -> &InMemoryPoisonStore {
51        &self.poison_manager
52    }
53}
54
55#[async_trait]
56impl ContextProvider for MockContextProvider {
57    async fn procedure_state(&self, procedure_id: ProcedureId) -> Result<Option<ProcedureState>> {
58        Ok(self.states.get(&procedure_id).cloned())
59    }
60
61    async fn procedure_state_receiver(
62        &self,
63        _procedure_id: ProcedureId,
64    ) -> Result<Option<Receiver<ProcedureState>>> {
65        Ok(None)
66    }
67
68    async fn try_put_poison(&self, key: &PoisonKey, procedure_id: ProcedureId) -> Result<()> {
69        self.poison_manager
70            .try_put_poison(key.to_string(), procedure_id.to_string())
71            .await
72    }
73
74    async fn acquire_lock(&self, key: &StringKey) -> DynamicKeyLockGuard {
75        acquire_dynamic_key_lock(&self.dynamic_key_lock, key).await
76    }
77}
78
79/// Executes a procedure until it returns [Status::Done].
80///
81/// # Panics
82/// Panics if the `procedure` has subprocedure to execute.
83pub async fn execute_procedure_until_done(procedure: &mut dyn Procedure) -> Option<Output> {
84    let ctx = Context {
85        procedure_id: ProcedureId::random(),
86        provider: Arc::new(MockContextProvider::default()),
87        event_context: None,
88    };
89
90    loop {
91        match procedure.execute(&ctx).await.unwrap() {
92            Status::Executing { .. } => (),
93            Status::Suspended { subprocedures, .. } => assert!(
94                subprocedures.is_empty(),
95                "Executing subprocedure is unsupported"
96            ),
97            Status::Done { output } => return output,
98            Status::Poisoned { .. } => return None,
99        }
100    }
101}
102
103/// Executes a procedure once.
104///
105/// Returns whether the procedure is done.
106pub async fn execute_procedure_once(
107    procedure_id: ProcedureId,
108    provider: MockContextProvider,
109    procedure: &mut dyn Procedure,
110) -> bool {
111    let ctx = Context {
112        procedure_id,
113        provider: Arc::new(provider),
114        event_context: None,
115    };
116
117    match procedure.execute(&ctx).await.unwrap() {
118        Status::Executing { .. } => false,
119        Status::Suspended { subprocedures, .. } => {
120            assert!(
121                subprocedures.is_empty(),
122                "Executing subprocedure is unsupported"
123            );
124            false
125        }
126        Status::Done { .. } => true,
127        Status::Poisoned { .. } => false,
128    }
129}
130
131/// Executes a procedure until it returns [Status::Suspended] or [Status::Done].
132///
133/// Returns `Some` if it returns [Status::Suspended] or `None` if it returns [Status::Done].
134pub async fn execute_until_suspended_or_done(
135    procedure_id: ProcedureId,
136    provider: MockContextProvider,
137    procedure: &mut dyn Procedure,
138) -> Option<Vec<ProcedureWithId>> {
139    let ctx = Context {
140        procedure_id,
141        provider: Arc::new(provider),
142        event_context: None,
143    };
144
145    loop {
146        match procedure.execute(&ctx).await.unwrap() {
147            Status::Executing { .. } => (),
148            Status::Suspended { subprocedures, .. } => return Some(subprocedures),
149            Status::Done { .. } => break,
150            Status::Poisoned { .. } => unreachable!(),
151        }
152    }
153
154    None
155}
156
157pub fn new_test_procedure_context() -> Context {
158    Context {
159        procedure_id: ProcedureId::random(),
160        provider: Arc::new(MockContextProvider::default()),
161        event_context: None,
162    }
163}
164
165pub async fn execute_procedure_until<P: Procedure>(procedure: &mut P, until: impl Fn(&P) -> bool) {
166    let mut reached = false;
167    let context = new_test_procedure_context();
168    while !matches!(
169        procedure.execute(&context).await.unwrap(),
170        Status::Done { .. }
171    ) {
172        if until(procedure) {
173            reached = true;
174            break;
175        }
176    }
177    assert!(
178        reached,
179        "procedure '{}' did not reach the expected state",
180        procedure.type_name()
181    );
182}
common_procedure_test/lib.rs

common_procedure_test/
lib.rs
