use std::{
collections::VecDeque,
fmt,
sync::{Arc, Mutex},
};
use tracing_core::{
span::{Attributes, Id, Record},
Event, Subscriber,
};
use tracing_subscriber::{
layer::{Context, Layer},
registry::{LookupSpan, SpanRef},
};
use crate::{
ancestry::{get_ancestry, ActualAncestry, HasAncestry},
event::ExpectedEvent,
expect::Expect,
span::{ActualSpan, ExpectedSpan, NewSpan},
subscriber::MockHandle,
};
#[must_use]
pub fn mock() -> MockLayerBuilder {
MockLayerBuilder {
expected: Default::default(),
name: std::thread::current()
.name()
.map(String::from)
.unwrap_or_default(),
}
}
#[must_use]
pub fn named(name: impl std::fmt::Display) -> MockLayerBuilder {
mock().named(name)
}
#[derive(Debug)]
pub struct MockLayerBuilder {
expected: VecDeque<Expect>,
name: String,
}
pub struct MockLayer {
expected: Arc<Mutex<VecDeque<Expect>>>,
current: Mutex<Vec<Id>>,
name: String,
}
impl MockLayerBuilder {
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 event(mut self, event: ExpectedEvent) -> Self {
self.expected.push_back(Expect::Event(event));
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 enter<S>(mut self, span: S) -> Self
where
S: Into<ExpectedSpan>,
{
self.expected.push_back(Expect::Enter(span.into()));
self
}
pub fn exit<S>(mut self, span: S) -> Self
where
S: Into<ExpectedSpan>,
{
self.expected.push_back(Expect::Exit(span.into()));
self
}
pub fn on_register_dispatch(mut self) -> Self {
self.expected.push_back(Expect::OnRegisterDispatch);
self
}
pub fn only(mut self) -> Self {
self.expected.push_back(Expect::Nothing);
self
}
pub fn run(self) -> MockLayer {
MockLayer {
expected: Arc::new(Mutex::new(self.expected)),
name: self.name,
current: Mutex::new(Vec::new()),
}
}
pub fn run_with_handle(self) -> (MockLayer, MockHandle) {
let expected = Arc::new(Mutex::new(self.expected));
let handle = MockHandle::new(expected.clone(), self.name.clone());
let subscriber = MockLayer {
expected,
name: self.name,
current: Mutex::new(Vec::new()),
};
(subscriber, handle)
}
}
impl<'a, S> From<&SpanRef<'a, S>> for ActualSpan
where
S: LookupSpan<'a>,
{
fn from(span_ref: &SpanRef<'a, S>) -> Self {
Self::new(span_ref.id(), Some(span_ref.metadata()))
}
}
impl MockLayer {
fn check_event_scope<C>(
&self,
current_scope: Option<tracing_subscriber::registry::Scope<'_, C>>,
expected_scope: &mut [ExpectedSpan],
) where
C: for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
{
let mut current_scope = current_scope.into_iter().flatten();
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
);
expected.check(
&(&actual).into(),
format_args!("the {}th span in the event's scope to be", i),
&self.name,
);
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,
);
}
}
impl<C> Layer<C> for MockLayer
where
C: Subscriber + for<'a> LookupSpan<'a>,
{
fn on_register_dispatch(&self, _subscriber: &tracing::Dispatch) {
println!("[{}] on_register_dispatch", self.name);
let mut expected = self.expected.lock().unwrap();
if let Some(Expect::OnRegisterDispatch) = expected.front() {
expected.pop_front();
}
}
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<'_, C>) {
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<'_, C>) {
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)) => {
expected.check(event, || context_get_ancestry(event, &cx), &self.name);
if let Some(expected_scope) = expected.scope_mut() {
self.check_event_scope(cx.event_scope(event), expected_scope);
}
}
Some(ex) => ex.bad(&self.name, format_args!("observed event {:#?}", event)),
}
}
fn on_follows_from(&self, _span: &Id, _follows: &Id, _: Context<'_, C>) {
unimplemented!(
"so far, we don't have any tests that need an `on_follows_from` \
implementation.\nif you just wrote one that does, feel free to \
implement it!"
);
}
fn on_new_span(&self, span: &Attributes<'_>, id: &Id, cx: Context<'_, C>) {
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() {
expected.check(span, || context_get_ancestry(span, &cx), &self.name);
}
}
}
fn on_enter(&self, id: &Id, cx: Context<'_, C>) {
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)) => {
expected_span.check(&(&span).into(), "to enter", &self.name);
}
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<'_, C>) {
if std::thread::panicking() {
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)) => {
expected_span.check(&(&span).into(), "to exit", &self.name);
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<'_, C>) {
if std::thread::panicking() {
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)) => {
if !::std::thread::panicking() {
if let Some(ref span) = span {
expected_span.check(&span.into(), "to close a span", &self.name);
}
}
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<'_, C>) {
panic!("well-behaved subscribers should never do this to us, lol");
}
}
fn context_get_ancestry<C>(item: impl HasAncestry, ctx: &Context<'_, C>) -> ActualAncestry
where
C: Subscriber + for<'a> LookupSpan<'a>,
{
get_ancestry(
item,
|| ctx.lookup_current().map(|s| s.id()),
|span_id| ctx.span(span_id).map(|span| (&span).into()),
)
}
impl fmt::Debug for MockLayer {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut s = f.debug_struct("ExpectSubscriber");
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!("{:?}", ¤t));
} else {
s.field("current", &format_args!("<locked>"));
}
s.finish()
}
}