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

aws_smithy_runtime/test_util/
capture_test_logs.rs
