#![allow(missing_docs)] use super::{ event::MockEvent, field as mock_field, span::{MockSpan, NewSpan}, }; use std::{ collections::{HashMap, VecDeque}, fmt, sync::{ atomic::{AtomicUsize, Ordering}, Arc, Mutex, }, }; use tokio_trace::{field, Event, Id, Metadata, Subscriber}; #[derive(Debug, Eq, PartialEq)] enum Expect { Event(MockEvent), Enter(MockSpan), Exit(MockSpan), CloneSpan(MockSpan), DropSpan(MockSpan), Record(MockSpan, mock_field::Expect), NewSpan(NewSpan), Nothing, } struct SpanState { name: &'static str, refs: usize, } struct Running bool> { spans: Mutex>, expected: Arc>>, ids: AtomicUsize, filter: F, } pub struct MockSubscriber bool> { expected: VecDeque, filter: F, } pub struct MockHandle(Arc>>); pub fn mock() -> MockSubscriber bool> { MockSubscriber { expected: VecDeque::new(), filter: (|_: &Metadata| true) as for<'r, 's> fn(&'r Metadata<'s>) -> _, } } impl bool> MockSubscriber { 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 clone_span(mut self, span: MockSpan) -> Self { self.expected.push_back(Expect::CloneSpan(span)); self } pub fn drop_span(mut self, span: MockSpan) -> Self { self.expected.push_back(Expect::DropSpan(span)); self } pub fn done(mut self) -> Self { self.expected.push_back(Expect::Nothing); self } pub fn record(mut self, span: MockSpan, fields: I) -> Self where I: Into, { self.expected.push_back(Expect::Record(span, fields.into())); self } pub fn new_span(mut self, new_span: I) -> Self where I: Into, { self.expected.push_back(Expect::NewSpan(new_span.into())); self } pub fn with_filter(self, filter: G) -> MockSubscriber where G: Fn(&Metadata) -> bool, { MockSubscriber { filter, expected: self.expected, } } 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()); let subscriber = Running { spans: Mutex::new(HashMap::new()), expected, ids: AtomicUsize::new(0), filter: self.filter, }; (subscriber, handle) } } impl bool> Subscriber for Running { fn enabled(&self, meta: &Metadata) -> bool { (self.filter)(meta) } fn record(&self, id: &Id, values: &field::ValueSet) { 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 {:?}", id)); println!("record: {}; id={:?}; values={:?};", span.name, id, values); let was_expected = if let Some(Expect::Record(_, _)) = expected.front() { true } else { false }; if was_expected { if let Expect::Record(expected_span, mut expected_values) = expected.pop_front().unwrap() { if let Some(name) = expected_span.name() { assert_eq!(name, span.name); } let mut checker = expected_values.checker(format!("span {}: ", span.name)); values.record(&mut checker); checker.finish(); } } } fn event(&self, event: &Event) { let name = event.metadata().name(); println!("event: {};", name); match self.expected.lock().unwrap().pop_front() { None => {} Some(Expect::Event(expected)) => expected.check(event), Some(ex) => ex.bad(format_args!("observed event {:?}", event)), } } fn record_follows_from(&self, _span: &Id, _follows: &Id) { // TODO: it should be possible to expect spans to follow from other spans } fn new_span(&self, meta: &Metadata, values: &field::ValueSet) -> Id { let id = self.ids.fetch_add(1, Ordering::SeqCst); let id = Id::from_u64(id as u64); println!( "new_span: name={:?}; target={:?}; id={:?};", meta.name(), meta.target(), id ); let mut expected = self.expected.lock().unwrap(); let was_expected = match expected.front() { Some(Expect::NewSpan(_)) => true, _ => false, }; if was_expected { if let Expect::NewSpan(mut expected) = expected.pop_front().unwrap() { let name = meta.name(); expected .span .metadata .check(meta, format_args!("span `{}`", name)); let mut checker = expected.fields.checker(format!("{}", name)); values.record(&mut checker); checker.finish(); } } self.spans.lock().unwrap().insert( id.clone(), SpanState { name: meta.name(), refs: 1, }, ); id } fn enter(&self, id: &Id) { let spans = self.spans.lock().unwrap(); if let Some(span) = spans.get(id) { println!("enter: {}; id={:?};", span.name, id); match self.expected.lock().unwrap().pop_front() { None => {} Some(Expect::Enter(ref expected_span)) => { if let Some(name) = expected_span.name() { assert_eq!(name, span.name); } } Some(ex) => ex.bad(format_args!("entered span {:?}", span.name)), } }; } fn exit(&self, id: &Id) { let spans = self.spans.lock().unwrap(); let span = spans .get(id) .unwrap_or_else(|| panic!("no span for ID {:?}", id)); println!("exit: {}; id={:?};", span.name, id); match self.expected.lock().unwrap().pop_front() { None => {} Some(Expect::Exit(ref expected_span)) => { if let Some(name) = expected_span.name() { assert_eq!(name, span.name); } } Some(ex) => ex.bad(format_args!("exited span {:?}", span.name)), }; } fn clone_span(&self, id: &Id) -> Id { let name = self.spans.lock().unwrap().get_mut(id).map(|span| { let name = span.name; println!("clone_span: {}; id={:?}; refs={:?};", name, id, span.refs); span.refs += 1; name }); if name.is_none() { println!("clone_span: id={:?};", id); } let mut expected = self.expected.lock().unwrap(); let was_expected = if let Some(Expect::CloneSpan(ref span)) = expected.front() { assert_eq!(name, span.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={:?};", name, id, span.refs); span.refs -= 1; name }) } else { None }; if name.is_none() { println!("drop_span: id={:?}", id); } if let Ok(mut expected) = self.expected.try_lock() { let was_expected = match expected.front() { Some(Expect::DropSpan(ref span)) => { // Don't assert if this function was called while panicking, // as failing the assertion can cause a double panic. if !::std::thread::panicking() { assert_eq!(name, span.name()); } true } Some(Expect::Event(_)) => { if !::std::thread::panicking() { assert!(is_event); } true } _ => false, }; if was_expected { expected.pop_front(); } } } } impl MockHandle { pub fn assert_finished(&self) { if let Ok(ref expected) = self.0.lock() { assert!( !expected.iter().any(|thing| thing != &Expect::Nothing), "more notifications expected: {:?}", **expected ); } } } impl Expect { fn bad<'a>(&self, what: fmt::Arguments<'a>) { match self { Expect::Event(e) => panic!("expected event {}, but {} instead", e, what,), Expect::Enter(e) => panic!("expected to enter {} but {} instead", e, what,), Expect::Exit(e) => panic!("expected to exit {} but {} instead", e, what,), Expect::CloneSpan(e) => panic!("expected to clone {} but {} instead", e, what,), Expect::DropSpan(e) => panic!("expected to drop {} but {} instead", e, what,), Expect::Record(e, fields) => { panic!("expected {} to record {} but {} instead", e, fields, what,) } Expect::NewSpan(e) => panic!("expected {} but {} instead", e, what), Expect::Nothing => panic!("expected nothing else to happen, but {} instead", what,), } } }