tracing-subscriber/tests/support.rs

#![allow(missing_docs, dead_code)]
pub use tracing_mock::{event, field, span, subscriber};

use tracing_core::{
    span::{Attributes, Id, Record},
    Event, Subscriber,
};
use tracing_mock::{
    event::MockEvent,
    span::{MockSpan, NewSpan},
    subscriber::{Expect, MockHandle},
};
use tracing_subscriber::{
    layer::{Context, Layer},
    registry::{LookupSpan, SpanRef},
};

use std::{
    collections::VecDeque,
    fmt,
    sync::{Arc, Mutex},
};

pub mod layer {
    use super::ExpectLayerBuilder;

    pub fn mock() -> ExpectLayerBuilder {
        ExpectLayerBuilder {
            expected: Default::default(),
            name: std::thread::current()
                .name()
                .map(String::from)
                .unwrap_or_default(),
        }
    }

    pub fn named(name: impl std::fmt::Display) -> ExpectLayerBuilder {
        mock().named(name)
    }
}

pub struct ExpectLayerBuilder {
    expected: VecDeque<Expect>,
    name: String,
}

pub struct ExpectLayer {
    expected: Arc<Mutex<VecDeque<Expect>>>,
    current: Mutex<Vec<Id>>,
    name: String,
}

impl ExpectLayerBuilder {
    /// Overrides the name printed by the mock subscriber's debugging output.
    ///
    /// The debugging output is displayed if the test panics, or if the test is
    /// run with `--nocapture`.
    ///
    /// By default, the mock subscriber's name is the  name of the test
    /// (*technically*, the name of the thread where it was created, which is
    /// the name of the test unless tests are run with `--test-threads=1`).
    /// When a test has only one mock subscriber, this is sufficient. However,
    /// some tests may include multiple subscribers, in order to test
    /// interactions between multiple subscribers. In that case, it can be
    /// helpful to give each subscriber a separate name to distinguish where the
    /// debugging output comes from.
    pub fn named(mut self, name: impl fmt::Display) -> Self {
        use std::fmt::Write;
        if !self.name.is_empty() {
            write!(&mut self.name, "::{}", name).unwrap();
        } else {
            self.name = name.to_string();
        }
        self
    }

    pub fn enter(mut self, span: MockSpan) -> Self {
        self.expected.push_back(Expect::Enter(span));
        self
    }

    pub fn event(mut self, event: MockEvent) -> Self {
        self.expected.push_back(Expect::Event(event));
        self
    }

    pub fn exit(mut self, span: MockSpan) -> Self {
        self.expected.push_back(Expect::Exit(span));
        self
    }

    pub fn done(mut self) -> Self {
        self.expected.push_back(Expect::Nothing);
        self
    }

    pub fn record<I>(mut self, span: MockSpan, fields: I) -> Self
    where
        I: Into<field::Expect>,
    {
        self.expected.push_back(Expect::Visit(span, fields.into()));
        self
    }

    pub fn new_span<I>(mut self, new_span: I) -> Self
    where
        I: Into<NewSpan>,
    {
        self.expected.push_back(Expect::NewSpan(new_span.into()));
        self
    }

    pub fn run(self) -> ExpectLayer {
        ExpectLayer {
            expected: Arc::new(Mutex::new(self.expected)),
            name: self.name,
            current: Mutex::new(Vec::new()),
        }
    }

    pub fn run_with_handle(self) -> (ExpectLayer, MockHandle) {
        let expected = Arc::new(Mutex::new(self.expected));
        let handle = MockHandle::new(expected.clone(), self.name.clone());
        let layer = ExpectLayer {
            expected,
            name: self.name,
            current: Mutex::new(Vec::new()),
        };
        (layer, handle)
    }
}

impl ExpectLayer {
    fn check_span_ref<'spans, S>(
        &self,
        expected: &MockSpan,
        actual: &SpanRef<'spans, S>,
        what_happened: impl fmt::Display,
    ) where
        S: LookupSpan<'spans>,
    {
        if let Some(exp_name) = expected.name() {
            assert_eq!(
                actual.name(),
                exp_name,
                "\n[{}] expected {} a span named {:?}\n\
                 [{}] but it was named {:?} instead (span {} {:?})",
                self.name,
                what_happened,
                exp_name,
                self.name,
                actual.name(),
                actual.name(),
                actual.id()
            );
        }

        if let Some(exp_level) = expected.level() {
            let actual_level = actual.metadata().level();
            assert_eq!(
                actual_level,
                &exp_level,
                "\n[{}] expected {} a span at {:?}\n\
                 [{}] but it was at {:?} instead (span {} {:?})",
                self.name,
                what_happened,
                exp_level,
                self.name,
                actual_level,
                actual.name(),
                actual.id(),
            );
        }

        if let Some(exp_target) = expected.target() {
            let actual_target = actual.metadata().target();
            assert_eq!(
                actual_target,
                exp_target,
                "\n[{}] expected {} a span with target {:?}\n\
                 [{}] but it had the target {:?} instead (span {} {:?})",
                self.name,
                what_happened,
                exp_target,
                self.name,
                actual_target,
                actual.name(),
                actual.id(),
            );
        }
    }
}

impl<S> Layer<S> for ExpectLayer
where
    S: Subscriber + for<'a> LookupSpan<'a>,
{
    fn register_callsite(
        &self,
        metadata: &'static tracing::Metadata<'static>,
    ) -> tracing_core::Interest {
        println!("[{}] register_callsite {:#?}", self.name, metadata);
        tracing_core::Interest::always()
    }

    fn on_record(&self, _: &Id, _: &Record<'_>, _: Context<'_, S>) {
        unimplemented!(
            "so far, we don't have any tests that need an `on_record` \
            implementation.\nif you just wrote one that does, feel free to \
            implement it!"
        );
    }

    fn on_event(&self, event: &Event<'_>, cx: Context<'_, S>) {
        let name = event.metadata().name();
        println!(
            "[{}] event: {}; level: {}; target: {}",
            self.name,
            name,
            event.metadata().level(),
            event.metadata().target(),
        );
        match self.expected.lock().unwrap().pop_front() {
            None => {}
            Some(Expect::Event(mut expected)) => {
                let get_parent_name = || cx.event_span(event).map(|span| span.name().to_string());
                expected.check(event, get_parent_name, &self.name);
                let mut current_scope = cx.event_scope(event).into_iter().flatten();
                let expected_scope = expected.scope_mut();
                let mut i = 0;
                for (expected, actual) in expected_scope.iter_mut().zip(&mut current_scope) {
                    println!(
                        "[{}] event_scope[{}] actual={} ({:?}); expected={}",
                        self.name,
                        i,
                        actual.name(),
                        actual.id(),
                        expected
                    );
                    self.check_span_ref(
                        expected,
                        &actual,
                        format_args!("the {}th span in the event's scope to be", i),
                    );
                    i += 1;
                }
                let remaining_expected = &expected_scope[i..];
                assert!(
                    remaining_expected.is_empty(),
                    "\n[{}] did not observe all expected spans in event scope!\n[{}] missing: {:#?}",
                    self.name,
                    self.name,
                    remaining_expected,
                );
                assert!(
                    current_scope.next().is_none(),
                    "\n[{}] did not expect all spans in the actual event scope!",
                    self.name,
                );
            }
            Some(ex) => ex.bad(&self.name, format_args!("observed event {:#?}", event)),
        }
    }

    fn on_follows_from(&self, _span: &Id, _follows: &Id, _: Context<'_, S>) {
        // TODO: it should be possible to expect spans to follow from other spans
    }

    fn on_new_span(&self, span: &Attributes<'_>, id: &Id, cx: Context<'_, S>) {
        let meta = span.metadata();
        println!(
            "[{}] new_span: name={:?}; target={:?}; id={:?};",
            self.name,
            meta.name(),
            meta.target(),
            id
        );
        let mut expected = self.expected.lock().unwrap();
        let was_expected = matches!(expected.front(), Some(Expect::NewSpan(_)));
        if was_expected {
            if let Expect::NewSpan(mut expected) = expected.pop_front().unwrap() {
                let get_parent_name = || {
                    span.parent()
                        .and_then(|id| cx.span(id))
                        .or_else(|| cx.lookup_current())
                        .map(|span| span.name().to_string())
                };
                expected.check(span, get_parent_name, &self.name);
            }
        }
    }

    fn on_enter(&self, id: &Id, cx: Context<'_, S>) {
        let span = cx
            .span(id)
            .unwrap_or_else(|| panic!("[{}] no span for ID {:?}", self.name, id));
        println!("[{}] enter: {}; id={:?};", self.name, span.name(), id);
        match self.expected.lock().unwrap().pop_front() {
            None => {}
            Some(Expect::Enter(ref expected_span)) => {
                self.check_span_ref(expected_span, &span, "to enter");
            }
            Some(ex) => ex.bad(&self.name, format_args!("entered span {:?}", span.name())),
        }
        self.current.lock().unwrap().push(id.clone());
    }

    fn on_exit(&self, id: &Id, cx: Context<'_, S>) {
        if std::thread::panicking() {
            // `exit()` can be called in `drop` impls, so we must guard against
            // double panics.
            println!("[{}] exit {:?} while panicking", self.name, id);
            return;
        }
        let span = cx
            .span(id)
            .unwrap_or_else(|| panic!("[{}] no span for ID {:?}", self.name, id));
        println!("[{}] exit: {}; id={:?};", self.name, span.name(), id);
        match self.expected.lock().unwrap().pop_front() {
            None => {}
            Some(Expect::Exit(ref expected_span)) => {
                self.check_span_ref(expected_span, &span, "to exit");
                let curr = self.current.lock().unwrap().pop();
                assert_eq!(
                    Some(id),
                    curr.as_ref(),
                    "[{}] exited span {:?}, but the current span was {:?}",
                    self.name,
                    span.name(),
                    curr.as_ref().and_then(|id| cx.span(id)).map(|s| s.name())
                );
            }
            Some(ex) => ex.bad(&self.name, format_args!("exited span {:?}", span.name())),
        };
    }

    fn on_close(&self, id: Id, cx: Context<'_, S>) {
        if std::thread::panicking() {
            // `try_close` can be called in `drop` impls, so we must guard against
            // double panics.
            println!("[{}] close {:?} while panicking", self.name, id);
            return;
        }
        let span = cx.span(&id);
        let name = span.as_ref().map(|span| {
            println!("[{}] close_span: {}; id={:?};", self.name, span.name(), id,);
            span.name()
        });
        if name.is_none() {
            println!("[{}] drop_span: id={:?}", self.name, id);
        }
        if let Ok(mut expected) = self.expected.try_lock() {
            let was_expected = match expected.front() {
                Some(Expect::DropSpan(ref expected_span)) => {
                    // Don't assert if this function was called while panicking,
                    // as failing the assertion can cause a double panic.
                    if !::std::thread::panicking() {
                        if let Some(ref span) = span {
                            self.check_span_ref(expected_span, span, "to close");
                        }
                    }
                    true
                }
                Some(Expect::Event(_)) => {
                    if !::std::thread::panicking() {
                        panic!(
                            "[{}] expected an event, but dropped span {} (id={:?}) instead",
                            self.name,
                            name.unwrap_or("<unknown name>"),
                            id
                        );
                    }
                    true
                }
                _ => false,
            };
            if was_expected {
                expected.pop_front();
            }
        }
    }

    fn on_id_change(&self, _old: &Id, _new: &Id, _ctx: Context<'_, S>) {
        panic!("well-behaved subscribers should never do this to us, lol");
    }
}

impl fmt::Debug for ExpectLayer {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut s = f.debug_struct("ExpectLayer");
        s.field("name", &self.name);

        if let Ok(expected) = self.expected.try_lock() {
            s.field("expected", &expected);
        } else {
            s.field("expected", &format_args!("<locked>"));
        }

        if let Ok(current) = self.current.try_lock() {
            s.field("current", &format_args!("{:?}", &current));
        } else {
            s.field("current", &format_args!("<locked>"));
        }

        s.finish()
    }
}