Skip to main content

diskann_tools/utils/
tracing.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use tracing;
7use tracing_subscriber::{filter::LevelFilter, fmt, prelude::*, EnvFilter};
8
9/// Create a default subscriber logging messages to `stdout` and respecting the `RUST_LOG`
10/// environment variable.
11///
12/// If the environment variable is not set - then the "info" level will be used.
13pub fn init_subscriber() {
14    let fmt_layer = fmt::layer().with_target(true);
15
16    let filter_layer = EnvFilter::builder()
17        .with_default_directive(LevelFilter::INFO.into())
18        .from_env_lossy();
19
20    tracing_subscriber::registry()
21        .with(filter_layer)
22        .with(fmt_layer)
23        .init();
24}
25
26/// Create a subscriber for the integration tests.
27///
28/// This subscriber returns a `Guard` that will only install the subscriber locally,
29/// allowing test threads to have non-conflicting subscribers.
30pub fn init_test_subscriber() -> tracing::subscriber::DefaultGuard {
31    let fmt_layer = fmt::layer().with_target(true).with_test_writer();
32
33    let filter_layer = EnvFilter::builder()
34        .with_default_directive(LevelFilter::INFO.into())
35        .from_env_lossy();
36
37    tracing_subscriber::registry()
38        .with(filter_layer)
39        .with(fmt_layer)
40        .set_default()
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46    use tracing::{debug, error, info, warn};
47
48    #[test]
49    fn test_init_test_subscriber() {
50        let _guard = init_test_subscriber();
51        // Test that logging works without panicking
52        info!("test info message");
53        warn!("test warn message");
54        error!("test error message");
55        debug!("test debug message");
56    }
57
58    #[test]
59    fn test_init_test_subscriber_guard_scope() {
60        {
61            let _guard = init_test_subscriber();
62            info!("inside guard scope");
63        }
64        // After guard is dropped, we can create a new one
65        let _guard2 = init_test_subscriber();
66        info!("new guard scope");
67    }
68}