aws_smithy_runtime/test_util/
capture_test_logs.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6use std::env;
7use std::io::Write;
8use std::sync::{Arc, Mutex};
9use tracing::subscriber::DefaultGuard;
10use tracing::Level;
11use tracing_subscriber::fmt::TestWriter;
12
13/// A guard that resets log capturing upon being dropped.
14#[derive(Debug)]
15pub struct LogCaptureGuard(#[allow(dead_code)] DefaultGuard);
16
17/// Enables output of test logs to stdout
18///
19/// The `env` filter can be configured to set the log level for different modules
20pub fn show_filtered_test_logs(filter: &str) -> LogCaptureGuard {
21    let (mut writer, _rx) = Tee::stdout();
22    writer.loud();
23
24    let subscriber = tracing_subscriber::fmt()
25        .with_env_filter(filter)
26        .with_writer(Mutex::new(writer))
27        .finish();
28    let guard = tracing::subscriber::set_default(subscriber);
29    LogCaptureGuard(guard)
30}
31
32/// Enables output of test logs to stdout at trace level by default.
33///
34/// The env filter can be changed with the `RUST_LOG` environment variable.
35#[must_use]
36pub fn show_test_logs() -> LogCaptureGuard {
37    let (mut writer, _rx) = Tee::stdout();
38    writer.loud();
39
40    let env_var = env::var("RUST_LOG").ok();
41    let env_filter = env_var.as_deref().unwrap_or("trace");
42    eprintln!(
43        "Enabled verbose test logging with env filter {env_filter:?}. \
44        You can change the env filter with the RUST_LOG environment variable."
45    );
46
47    show_filtered_test_logs(env_filter)
48}
49
50/// Capture logs from this test.
51///
52/// The logs will be captured until the `DefaultGuard` is dropped.
53///
54/// *Why use this instead of traced_test?*
55/// This captures _all_ logs, not just logs produced by the current crate.
56#[must_use] // log capturing ceases the instant the `DefaultGuard` is dropped
57pub fn capture_test_logs() -> (LogCaptureGuard, Rx) {
58    // it may be helpful to upstream this at some point
59    let (mut writer, rx) = Tee::stdout();
60    if env::var("VERBOSE_TEST_LOGS").is_ok() {
61        eprintln!("Enabled verbose test logging.");
62        writer.loud();
63    } else {
64        eprintln!("To see full logs from this test set VERBOSE_TEST_LOGS=true");
65    }
66    let subscriber = tracing_subscriber::fmt()
67        .with_max_level(Level::TRACE)
68        .with_writer(Mutex::new(writer))
69        .finish();
70    let guard = tracing::subscriber::set_default(subscriber);
71    (LogCaptureGuard(guard), rx)
72}
73
74/// Receiver for the captured logs.
75pub struct Rx(Arc<Mutex<Vec<u8>>>);
76impl Rx {
77    /// Returns the captured logs as a string.
78    ///
79    /// # Panics
80    /// This will panic if the logs are not valid UTF-8.
81    pub fn contents(&self) -> String {
82        String::from_utf8(self.0.lock().unwrap().clone()).unwrap()
83    }
84}
85
86struct Tee<W> {
87    buf: Arc<Mutex<Vec<u8>>>,
88    quiet: bool,
89    inner: W,
90}
91
92impl Tee<TestWriter> {
93    fn stdout() -> (Self, Rx) {
94        let buf: Arc<Mutex<Vec<u8>>> = Default::default();
95        (
96            Tee {
97                buf: buf.clone(),
98                quiet: true,
99                inner: TestWriter::new(),
100            },
101            Rx(buf),
102        )
103    }
104}
105
106impl<W> Tee<W> {
107    fn loud(&mut self) {
108        self.quiet = false;
109    }
110}
111
112impl<W> Write for Tee<W>
113where
114    W: Write,
115{
116    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
117        self.buf.lock().unwrap().extend_from_slice(buf);
118        if !self.quiet {
119            self.inner.write_all(buf)?;
120            Ok(buf.len())
121        } else {
122            Ok(buf.len())
123        }
124    }
125
126    fn flush(&mut self) -> std::io::Result<()> {
127        self.inner.flush()
128    }
129}
aws_smithy_runtime/test_util/capture_test_logs.rs

aws_smithy_runtime/test_util/
capture_test_logs.rs
