Skip to main content

lance_testing/
progress.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4/// Define a test-only progress recorder that captures stage events in memory.
5#[macro_export]
6macro_rules! define_stage_event_progress {
7    ($name:ident, $progress_trait:path, $result:ty) => {
8        #[derive(Debug, Default)]
9        struct $name {
10            events: ::std::sync::Mutex<Vec<(String, String, u64)>>,
11        }
12
13        impl $name {
14            fn recorded_events(&self) -> Vec<(String, String, u64)> {
15                self.events
16                    .lock()
17                    .expect("recording progress mutex poisoned")
18                    .clone()
19            }
20        }
21
22        #[::async_trait::async_trait]
23        impl $progress_trait for $name {
24            async fn stage_start(&self, stage: &str, total: Option<u64>, _unit: &str) -> $result {
25                self.events
26                    .lock()
27                    .expect("recording progress mutex poisoned")
28                    .push(("start".to_string(), stage.to_string(), total.unwrap_or(0)));
29                Ok(())
30            }
31
32            async fn stage_progress(&self, stage: &str, completed: u64) -> $result {
33                self.events
34                    .lock()
35                    .expect("recording progress mutex poisoned")
36                    .push(("progress".to_string(), stage.to_string(), completed));
37                Ok(())
38            }
39
40            async fn stage_complete(&self, stage: &str) -> $result {
41                self.events
42                    .lock()
43                    .expect("recording progress mutex poisoned")
44                    .push(("complete".to_string(), stage.to_string(), 0));
45                Ok(())
46            }
47        }
48    };
49}