use std::{
collections::{HashMap, VecDeque},
sync::{
atomic::{AtomicUsize, Ordering},
Arc, Mutex,
},
thread,
};
use tracing::{
level_filters::LevelFilter,
span::{self, Attributes, Id},
subscriber::Interest,
Event, Metadata, Subscriber,
};
use crate::{
ancestry::get_ancestry,
event::ExpectedEvent,
expect::Expect,
field::ExpectedFields,
span::{ActualSpan, ExpectedSpan, NewSpan},
};
pub(crate) struct SpanState {
id: Id,
name: &'static str,
refs: usize,
meta: &'static Metadata<'static>,
}
impl From<&SpanState> for ActualSpan {
fn from(span_state: &SpanState) -> Self {
Self::new(span_state.id.clone(), Some(span_state.meta))
}
}
struct Running<F: Fn(&Metadata<'_>) -> bool> {
spans: Mutex<HashMap<Id, SpanState>>,
expected: Arc<Mutex<VecDeque<Expect>>>,
current: Mutex<Vec<Id>>,
ids: AtomicUsize,
max_level: Option<LevelFilter>,
filter: F,
name: String,
}
#[derive(Debug)]
pub struct MockSubscriber<F: Fn(&Metadata<'_>) -> bool> {
expected: VecDeque<Expect>,
max_level: Option<LevelFilter>,
filter: F,
name: String,
}
#[derive(Debug)]
pub struct MockHandle(Arc<Mutex<VecDeque<Expect>>>, String);
#[must_use]
pub fn mock() -> MockSubscriber<fn(&Metadata<'_>) -> bool> {
MockSubscriber {
expected: VecDeque::new(),
filter: (|_: &Metadata<'_>| true) as for<'r, 's> fn(&'r Metadata<'s>) -> _,
max_level: None,
name: thread::current()
.name()
.unwrap_or("mock_subscriber")
.to_string(),
}
}
impl<F> MockSubscriber<F>
where
F: Fn(&Metadata<'_>) -> bool + 'static,
{
pub fn named(self, name: impl ToString) -> Self {
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 clone_span<S>(mut self, span: S) -> Self
where
S: Into<ExpectedSpan>,
{
self.expected.push_back(Expect::CloneSpan(span.into()));
self
}
#[allow(deprecated)]
pub fn drop_span<S>(mut self, span: S) -> Self
where
S: Into<ExpectedSpan>,
{
self.expected.push_back(Expect::DropSpan(span.into()));
self
}
pub fn follows_from<S1, S2>(mut self, consequence: S1, cause: S2) -> Self
where
S1: Into<ExpectedSpan>,
S2: Into<ExpectedSpan>,
{
self.expected.push_back(Expect::FollowsFrom {
consequence: consequence.into(),
cause: cause.into(),
});
self
}
pub fn record<S, I>(mut self, span: S, fields: I) -> Self
where
S: Into<ExpectedSpan>,
I: Into<ExpectedFields>,
{
self.expected
.push_back(Expect::Visit(span.into(), fields.into()));
self
}
pub fn on_register_dispatch(mut self) -> Self {
self.expected.push_back(Expect::OnRegisterDispatch);
self
}
pub fn with_filter<G>(self, filter: G) -> MockSubscriber<G>
where
G: Fn(&Metadata<'_>) -> bool + 'static,
{
MockSubscriber {
expected: self.expected,
filter,
max_level: self.max_level,
name: self.name,
}
}
pub fn with_max_level_hint(self, hint: impl Into<LevelFilter>) -> Self {
Self {
max_level: Some(hint.into()),
..self
}
}
pub fn only(mut self) -> Self {
self.expected.push_back(Expect::Nothing);
self
}
pub fn run(self) -> impl Subscriber {
let (subscriber, _) = self.run_with_handle();
subscriber
}
pub fn run_with_handle(self) -> (impl Subscriber, MockHandle) {
let expected = Arc::new(Mutex::new(self.expected));
let handle = MockHandle(expected.clone(), self.name.clone());
let subscriber = Running {
spans: Mutex::new(HashMap::new()),
expected,
current: Mutex::new(Vec::new()),
ids: AtomicUsize::new(1),
filter: self.filter,
max_level: self.max_level,
name: self.name,
};
(subscriber, handle)
}
}
impl<F> Subscriber for Running<F>
where
F: Fn(&Metadata<'_>) -> bool + 'static,
{
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 enabled(&self, meta: &Metadata<'_>) -> bool {
println!("[{}] enabled: {:#?}", self.name, meta);
let enabled = (self.filter)(meta);
println!("[{}] enabled -> {}", self.name, enabled);
enabled
}
fn register_callsite(&self, meta: &'static Metadata<'static>) -> Interest {
println!("[{}] register_callsite: {:#?}", self.name, meta);
if self.enabled(meta) {
Interest::always()
} else {
Interest::never()
}
}
fn max_level_hint(&self) -> Option<LevelFilter> {
self.max_level
}
fn record(&self, id: &Id, values: &span::Record<'_>) {
let spans = self.spans.lock().unwrap();
let mut expected = self.expected.lock().unwrap();
let span = spans
.get(id)
.unwrap_or_else(|| panic!("[{}] no span for ID {:?}", self.name, id));
println!(
"[{}] record: {}; id={:?}; values={:?};",
self.name, span.name, id, values
);
let was_expected = matches!(expected.front(), Some(Expect::Visit(_, _)));
if was_expected {
if let Expect::Visit(expected_span, mut expected_values) = expected.pop_front().unwrap()
{
if let Some(name) = expected_span.name() {
assert_eq!(name, span.name);
}
let context = format!("span {}: ", span.name);
let mut checker = expected_values.checker(&context, &self.name);
values.record(&mut checker);
checker.finish();
}
}
}
fn event(&self, event: &Event<'_>) {
let name = event.metadata().name();
println!("[{}] event: {};", self.name, name);
match self.expected.lock().unwrap().pop_front() {
None => {}
Some(Expect::Event(mut expected)) => {
#[cfg(feature = "tracing-subscriber")]
{
if expected.scope_mut().is_some() {
unimplemented!(
"Expected scope for events is not supported with `MockSubscriber`."
)
}
}
let event_get_ancestry = || {
get_ancestry(
event,
|| self.lookup_current(),
|span_id| {
self.spans
.lock()
.unwrap()
.get(span_id)
.map(|span| span.into())
},
)
};
expected.check(event, event_get_ancestry, &self.name);
}
Some(ex) => ex.bad(&self.name, format_args!("observed event {:#?}", event)),
}
}
fn record_follows_from(&self, consequence_id: &Id, cause_id: &Id) {
let spans = self.spans.lock().unwrap();
if let Some(consequence_span) = spans.get(consequence_id) {
if let Some(cause_span) = spans.get(cause_id) {
println!(
"[{}] record_follows_from: {} (id={:?}) follows {} (id={:?})",
self.name, consequence_span.name, consequence_id, cause_span.name, cause_id,
);
match self.expected.lock().unwrap().pop_front() {
None => {}
Some(Expect::FollowsFrom {
consequence: ref expected_consequence,
cause: ref expected_cause,
}) => {
if let Some(name) = expected_consequence.name() {
assert_eq!(name, consequence_span.name);
}
if let Some(name) = expected_cause.name() {
assert_eq!(name, cause_span.name);
}
}
Some(ex) => ex.bad(
&self.name,
format_args!(
"consequence {:?} followed cause {:?}",
consequence_span.name, cause_span.name
),
),
}
}
};
}
fn new_span(&self, span: &Attributes<'_>) -> Id {
let meta = span.metadata();
let id = self.ids.fetch_add(1, Ordering::SeqCst);
let id = Id::from_u64(id as u64);
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(_)));
let mut spans = self.spans.lock().unwrap();
if was_expected {
if let Expect::NewSpan(mut expected) = expected.pop_front().unwrap() {
if let Some(expected_id) = &expected.span.id {
expected_id.set(id.into_u64()).unwrap();
}
expected.check(
span,
|| {
get_ancestry(
span,
|| self.lookup_current(),
|span_id| spans.get(span_id).map(|span| span.into()),
)
},
&self.name,
);
}
}
spans.insert(
id.clone(),
SpanState {
id: id.clone(),
name: meta.name(),
refs: 1,
meta,
},
);
id
}
fn enter(&self, id: &Id) {
let spans = self.spans.lock().unwrap();
if let Some(span) = spans.get(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 a span", &self.name);
}
Some(ex) => ex.bad(&self.name, format_args!("entered span {:?}", span.name)),
}
};
self.current.lock().unwrap().push(id.clone());
}
fn exit(&self, id: &Id) {
if std::thread::panicking() {
println!("[{}] exit {:?} while panicking", self.name, id);
return;
}
let spans = self.spans.lock().unwrap();
let span = spans
.get(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 a span", &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| spans.get(id)).map(|s| s.name)
);
}
Some(ex) => ex.bad(&self.name, format_args!("exited span {:?}", span.name)),
};
}
fn clone_span(&self, id: &Id) -> Id {
let mut spans = self.spans.lock().unwrap();
let mut span = spans.get_mut(id);
match span.as_deref_mut() {
Some(span) => {
println!(
"[{}] clone_span: {}; id={:?}; refs={:?};",
self.name, span.name, id, span.refs,
);
span.refs += 1;
}
None => {
println!(
"[{}] clone_span: id={:?} (not found in span list);",
self.name, id
);
}
}
let mut expected = self.expected.lock().unwrap();
let was_expected = if let Some(Expect::CloneSpan(ref expected_span)) = expected.front() {
match span {
Some(actual_span) => {
let actual_span: &_ = actual_span;
expected_span.check(&actual_span.into(), "to clone a span", &self.name);
}
None => expected_span.check(&id.into(), "to clone a span", &self.name),
}
true
} else {
false
};
if was_expected {
expected.pop_front();
}
id.clone()
}
fn drop_span(&self, id: Id) {
let mut is_event = false;
let name = if let Ok(mut spans) = self.spans.try_lock() {
spans.get_mut(&id).map(|span| {
let name = span.name;
if name.contains("event") {
is_event = true;
}
println!(
"[{}] drop_span: {}; id={:?}; refs={:?};",
self.name, name, id, span.refs
);
span.refs -= 1;
name
})
} else {
None
};
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 span)) => {
if !::std::thread::panicking() {
assert_eq!(name, span.name());
}
true
}
Some(Expect::Event(_)) => {
if !::std::thread::panicking() {
assert!(is_event, "[{}] expected an event", self.name);
}
true
}
_ => false,
};
if was_expected {
expected.pop_front();
}
}
}
fn current_span(&self) -> tracing_core::span::Current {
let stack = self.current.lock().unwrap();
match stack.last() {
Some(id) => {
let spans = self.spans.lock().unwrap();
let state = spans.get(id).expect("state for current span");
tracing_core::span::Current::new(id.clone(), state.meta)
}
None => tracing_core::span::Current::none(),
}
}
}
impl<F> Running<F>
where
F: Fn(&Metadata<'_>) -> bool,
{
fn lookup_current(&self) -> Option<span::Id> {
let stack = self.current.lock().unwrap();
stack.last().cloned()
}
}
impl MockHandle {
#[cfg(feature = "tracing-subscriber")]
pub(crate) fn new(expected: Arc<Mutex<VecDeque<Expect>>>, name: String) -> Self {
Self(expected, name)
}
pub fn assert_finished(&self) {
if let Ok(ref expected) = self.0.lock() {
assert!(
!expected.iter().any(|thing| thing != &Expect::Nothing),
"\n[{}] more notifications expected: {:#?}",
self.1,
**expected
);
}
}
}