diskann_benchmark_runner/
output.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6/// Return the default implementation of [`Output`] that sends:
7///
8/// * Prints to `stdout`.
9/// * Progress Bars to `stderr`.
10pub fn default() -> DefaultOutput {
11    DefaultOutput::new()
12}
13
14/// To enable testing and output redirection, the output stream that tests print to is hidden
15/// behind this `Output` trait, which consists of two parts:
16///
17/// * `sink`: A `&mut dyn std::io::Write` which is the target for all test prints.
18/// * `draw_target`: The target that all progress bars should use.
19pub trait Output {
20    fn sink(&mut self) -> &mut dyn std::io::Write;
21    fn draw_target(&self) -> indicatif::ProgressDrawTarget;
22}
23
24/// This allows `&mut dyn Output` to be used as the receiver of the `write!` macro.
25impl std::io::Write for &mut dyn Output {
26    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
27        self.sink().write(buf)
28    }
29    fn flush(&mut self) -> std::io::Result<()> {
30        self.sink().flush()
31    }
32}
33
34/// A default output that sends:
35///
36/// * Prints to `stdout`
37/// * Progress Bars to `stderr`.
38#[derive(Debug)]
39pub struct DefaultOutput(std::io::Stdout);
40
41impl DefaultOutput {
42    /// Construct a new [`DefaultOutput`].
43    pub fn new() -> Self {
44        Self(std::io::stdout())
45    }
46}
47
48impl Default for DefaultOutput {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl Output for DefaultOutput {
55    fn sink(&mut self) -> &mut dyn std::io::Write {
56        &mut self.0
57    }
58
59    // Note: Progress bars use `stderr` as the default draw target.
60    fn draw_target(&self) -> indicatif::ProgressDrawTarget {
61        indicatif::ProgressDrawTarget::stderr()
62    }
63}
64
65/// An output that suppresses all prints and progress bars.
66#[derive(Debug)]
67pub struct Sink(std::io::Sink);
68
69impl Sink {
70    /// Construct a new [`Sink`] output.
71    pub fn new() -> Self {
72        Self(std::io::sink())
73    }
74}
75
76impl Default for Sink {
77    fn default() -> Self {
78        Self::new()
79    }
80}
81
82impl Output for Sink {
83    fn sink(&mut self) -> &mut dyn std::io::Write {
84        &mut self.0
85    }
86
87    // Note: Progress bars use `stderr` as the default draw target.
88    fn draw_target(&self) -> indicatif::ProgressDrawTarget {
89        indicatif::ProgressDrawTarget::hidden()
90    }
91}
92
93/// The `Memory` allows test output to be captured directly in a buffer.
94///
95/// Integration tests can use this so each tests runs in its own environment.
96///
97/// Progress bars are suppressed.
98#[derive(Debug)]
99pub struct Memory(Vec<u8>);
100
101impl Memory {
102    /// Construct a new [`Memory`] with an empty buffer.
103    pub fn new() -> Self {
104        Self(Vec::new())
105    }
106
107    /// Consume `self`, returning the interior buffer with all messages that have been
108    /// written to `self`.
109    pub fn into_inner(self) -> Vec<u8> {
110        self.0
111    }
112}
113
114impl Default for Memory {
115    fn default() -> Self {
116        Self::new()
117    }
118}
119
120impl Output for Memory {
121    fn sink(&mut self) -> &mut dyn std::io::Write {
122        &mut self.0
123    }
124
125    // Hide the progress bar when piping to a buffer.
126    fn draw_target(&self) -> indicatif::ProgressDrawTarget {
127        indicatif::ProgressDrawTarget::hidden()
128    }
129}
130
131///////////
132// Tests //
133///////////
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    use std::io::Write;
140
141    #[test]
142    fn test_memory() {
143        let mut buf = Memory::new();
144        {
145            let mut output: &mut dyn Output = &mut buf;
146            writeln!(output, "hello world").unwrap();
147            writeln!(output, "test: {}", 10).unwrap();
148            output.flush().unwrap();
149
150            assert!(output.draw_target().is_hidden());
151        }
152        let bytes = buf.into_inner();
153        let message = str::from_utf8(&bytes).unwrap();
154        let mut lines = message.lines();
155        let first = lines.next().unwrap();
156        assert_eq!(first, "hello world");
157
158        let second = lines.next().unwrap();
159        assert_eq!(second, "test: 10");
160
161        assert!(lines.next().is_none());
162    }
163
164    #[test]
165    fn test_default() {
166        let mut d = default();
167        let mut s = &mut d as &mut dyn Output;
168
169        // This is not a reliable test when running under tooling like LLVM-cov.
170        // assert!(!s.draw_target().is_hidden());
171
172        writeln!(s, "test").unwrap();
173    }
174
175    #[test]
176    fn test_sink() {
177        let mut d = Sink::new();
178        let mut s = &mut d as &mut dyn Output;
179
180        assert!(s.draw_target().is_hidden());
181        writeln!(s, "test").unwrap();
182    }
183}
diskann_benchmark_runner/output.rs

diskann_benchmark_runner/
output.rs
