1 #![allow(missing_docs)]
2 use super::{
3 event::MockEvent,
4 field as mock_field,
5 span::{MockSpan, NewSpan},
6 };
7 use std::{
8 collections::{HashMap, VecDeque},
9 fmt,
10 sync::{
11 atomic::{AtomicUsize, Ordering},
12 Arc, Mutex,
13 },
14 thread,
15 };
16 use tracing::{
17 level_filters::LevelFilter,
18 span::{self, Attributes, Id},
19 subscriber::Interest,
20 Event, Metadata, Subscriber,
21 };
22
23 #[derive(Debug, Eq, PartialEq)]
24 pub enum Expect {
25 Event(MockEvent),
26 Enter(MockSpan),
27 Exit(MockSpan),
28 CloneSpan(MockSpan),
29 DropSpan(MockSpan),
30 Visit(MockSpan, mock_field::Expect),
31 NewSpan(NewSpan),
32 Nothing,
33 }
34
35 struct SpanState {
36 name: &'static str,
37 refs: usize,
38 }
39
40 struct Running<F: Fn(&Metadata<'_>) -> bool> {
41 spans: Mutex<HashMap<Id, SpanState>>,
42 expected: Arc<Mutex<VecDeque<Expect>>>,
43 current: Mutex<Vec<Id>>,
44 ids: AtomicUsize,
45 max_level: Option<LevelFilter>,
46 filter: F,
47 name: String,
48 }
49
50 pub struct MockSubscriber<F: Fn(&Metadata<'_>) -> bool> {
51 expected: VecDeque<Expect>,
52 max_level: Option<LevelFilter>,
53 filter: F,
54 name: String,
55 }
56
57 pub struct MockHandle(Arc<Mutex<VecDeque<Expect>>>, String);
58
mock() -> MockSubscriber<fn(&Metadata<'_>) -> bool>59 pub fn mock() -> MockSubscriber<fn(&Metadata<'_>) -> bool> {
60 MockSubscriber {
61 expected: VecDeque::new(),
62 filter: (|_: &Metadata<'_>| true) as for<'r, 's> fn(&'r Metadata<'s>) -> _,
63 max_level: None,
64 name: thread::current()
65 .name()
66 .unwrap_or("mock_subscriber")
67 .to_string(),
68 }
69 }
70
71 impl<F> MockSubscriber<F>
72 where
73 F: Fn(&Metadata<'_>) -> bool + 'static,
74 {
75 /// Overrides the name printed by the mock subscriber's debugging output.
76 ///
77 /// The debugging output is displayed if the test panics, or if the test is
78 /// run with `--nocapture`.
79 ///
80 /// By default, the mock subscriber's name is the name of the test
81 /// (*technically*, the name of the thread where it was created, which is
82 /// the name of the test unless tests are run with `--test-threads=1`).
83 /// When a test has only one mock subscriber, this is sufficient. However,
84 /// some tests may include multiple subscribers, in order to test
85 /// interactions between multiple subscribers. In that case, it can be
86 /// helpful to give each subscriber a separate name to distinguish where the
87 /// debugging output comes from.
named(self, name: impl ToString) -> Self88 pub fn named(self, name: impl ToString) -> Self {
89 Self {
90 name: name.to_string(),
91 ..self
92 }
93 }
94
enter(mut self, span: MockSpan) -> Self95 pub fn enter(mut self, span: MockSpan) -> Self {
96 self.expected.push_back(Expect::Enter(span));
97 self
98 }
99
event(mut self, event: MockEvent) -> Self100 pub fn event(mut self, event: MockEvent) -> Self {
101 self.expected.push_back(Expect::Event(event));
102 self
103 }
104
exit(mut self, span: MockSpan) -> Self105 pub fn exit(mut self, span: MockSpan) -> Self {
106 self.expected.push_back(Expect::Exit(span));
107 self
108 }
109
clone_span(mut self, span: MockSpan) -> Self110 pub fn clone_span(mut self, span: MockSpan) -> Self {
111 self.expected.push_back(Expect::CloneSpan(span));
112 self
113 }
114
115 #[allow(deprecated)]
drop_span(mut self, span: MockSpan) -> Self116 pub fn drop_span(mut self, span: MockSpan) -> Self {
117 self.expected.push_back(Expect::DropSpan(span));
118 self
119 }
120
done(mut self) -> Self121 pub fn done(mut self) -> Self {
122 self.expected.push_back(Expect::Nothing);
123 self
124 }
125
record<I>(mut self, span: MockSpan, fields: I) -> Self where I: Into<mock_field::Expect>,126 pub fn record<I>(mut self, span: MockSpan, fields: I) -> Self
127 where
128 I: Into<mock_field::Expect>,
129 {
130 self.expected.push_back(Expect::Visit(span, fields.into()));
131 self
132 }
133
new_span<I>(mut self, new_span: I) -> Self where I: Into<NewSpan>,134 pub fn new_span<I>(mut self, new_span: I) -> Self
135 where
136 I: Into<NewSpan>,
137 {
138 self.expected.push_back(Expect::NewSpan(new_span.into()));
139 self
140 }
141
with_filter<G>(self, filter: G) -> MockSubscriber<G> where G: Fn(&Metadata<'_>) -> bool + 'static,142 pub fn with_filter<G>(self, filter: G) -> MockSubscriber<G>
143 where
144 G: Fn(&Metadata<'_>) -> bool + 'static,
145 {
146 MockSubscriber {
147 expected: self.expected,
148 filter,
149 max_level: self.max_level,
150 name: self.name,
151 }
152 }
153
with_max_level_hint(self, hint: impl Into<LevelFilter>) -> Self154 pub fn with_max_level_hint(self, hint: impl Into<LevelFilter>) -> Self {
155 Self {
156 max_level: Some(hint.into()),
157 ..self
158 }
159 }
160
run(self) -> impl Subscriber161 pub fn run(self) -> impl Subscriber {
162 let (subscriber, _) = self.run_with_handle();
163 subscriber
164 }
165
run_with_handle(self) -> (impl Subscriber, MockHandle)166 pub fn run_with_handle(self) -> (impl Subscriber, MockHandle) {
167 let expected = Arc::new(Mutex::new(self.expected));
168 let handle = MockHandle(expected.clone(), self.name.clone());
169 let subscriber = Running {
170 spans: Mutex::new(HashMap::new()),
171 expected,
172 current: Mutex::new(Vec::new()),
173 ids: AtomicUsize::new(1),
174 filter: self.filter,
175 max_level: self.max_level,
176 name: self.name,
177 };
178 (subscriber, handle)
179 }
180 }
181
182 impl<F> Subscriber for Running<F>
183 where
184 F: Fn(&Metadata<'_>) -> bool + 'static,
185 {
enabled(&self, meta: &Metadata<'_>) -> bool186 fn enabled(&self, meta: &Metadata<'_>) -> bool {
187 println!("[{}] enabled: {:#?}", self.name, meta);
188 let enabled = (self.filter)(meta);
189 println!("[{}] enabled -> {}", self.name, enabled);
190 enabled
191 }
192
register_callsite(&self, meta: &'static Metadata<'static>) -> Interest193 fn register_callsite(&self, meta: &'static Metadata<'static>) -> Interest {
194 println!("[{}] register_callsite: {:#?}", self.name, meta);
195 if self.enabled(meta) {
196 Interest::always()
197 } else {
198 Interest::never()
199 }
200 }
max_level_hint(&self) -> Option<LevelFilter>201 fn max_level_hint(&self) -> Option<LevelFilter> {
202 self.max_level
203 }
204
record(&self, id: &Id, values: &span::Record<'_>)205 fn record(&self, id: &Id, values: &span::Record<'_>) {
206 let spans = self.spans.lock().unwrap();
207 let mut expected = self.expected.lock().unwrap();
208 let span = spans
209 .get(id)
210 .unwrap_or_else(|| panic!("[{}] no span for ID {:?}", self.name, id));
211 println!(
212 "[{}] record: {}; id={:?}; values={:?};",
213 self.name, span.name, id, values
214 );
215 let was_expected = matches!(expected.front(), Some(Expect::Visit(_, _)));
216 if was_expected {
217 if let Expect::Visit(expected_span, mut expected_values) = expected.pop_front().unwrap()
218 {
219 if let Some(name) = expected_span.name() {
220 assert_eq!(name, span.name);
221 }
222 let context = format!("span {}: ", span.name);
223 let mut checker = expected_values.checker(&context, &self.name);
224 values.record(&mut checker);
225 checker.finish();
226 }
227 }
228 }
229
event(&self, event: &Event<'_>)230 fn event(&self, event: &Event<'_>) {
231 let name = event.metadata().name();
232 println!("[{}] event: {};", self.name, name);
233 match self.expected.lock().unwrap().pop_front() {
234 None => {}
235 Some(Expect::Event(mut expected)) => {
236 let get_parent_name = || {
237 let stack = self.current.lock().unwrap();
238 let spans = self.spans.lock().unwrap();
239 event
240 .parent()
241 .and_then(|id| spans.get(id))
242 .or_else(|| stack.last().and_then(|id| spans.get(id)))
243 .map(|s| s.name.to_string())
244 };
245 expected.check(event, get_parent_name, &self.name);
246 }
247 Some(ex) => ex.bad(&self.name, format_args!("observed event {:#?}", event)),
248 }
249 }
250
record_follows_from(&self, _span: &Id, _follows: &Id)251 fn record_follows_from(&self, _span: &Id, _follows: &Id) {
252 // TODO: it should be possible to expect spans to follow from other spans
253 }
254
new_span(&self, span: &Attributes<'_>) -> Id255 fn new_span(&self, span: &Attributes<'_>) -> Id {
256 let meta = span.metadata();
257 let id = self.ids.fetch_add(1, Ordering::SeqCst);
258 let id = Id::from_u64(id as u64);
259 println!(
260 "[{}] new_span: name={:?}; target={:?}; id={:?};",
261 self.name,
262 meta.name(),
263 meta.target(),
264 id
265 );
266 let mut expected = self.expected.lock().unwrap();
267 let was_expected = matches!(expected.front(), Some(Expect::NewSpan(_)));
268 let mut spans = self.spans.lock().unwrap();
269 if was_expected {
270 if let Expect::NewSpan(mut expected) = expected.pop_front().unwrap() {
271 let get_parent_name = || {
272 let stack = self.current.lock().unwrap();
273 span.parent()
274 .and_then(|id| spans.get(id))
275 .or_else(|| stack.last().and_then(|id| spans.get(id)))
276 .map(|s| s.name.to_string())
277 };
278 expected.check(span, get_parent_name, &self.name);
279 }
280 }
281 spans.insert(
282 id.clone(),
283 SpanState {
284 name: meta.name(),
285 refs: 1,
286 },
287 );
288 id
289 }
290
enter(&self, id: &Id)291 fn enter(&self, id: &Id) {
292 let spans = self.spans.lock().unwrap();
293 if let Some(span) = spans.get(id) {
294 println!("[{}] enter: {}; id={:?};", self.name, span.name, id);
295 match self.expected.lock().unwrap().pop_front() {
296 None => {}
297 Some(Expect::Enter(ref expected_span)) => {
298 if let Some(name) = expected_span.name() {
299 assert_eq!(name, span.name);
300 }
301 }
302 Some(ex) => ex.bad(&self.name, format_args!("entered span {:?}", span.name)),
303 }
304 };
305 self.current.lock().unwrap().push(id.clone());
306 }
307
exit(&self, id: &Id)308 fn exit(&self, id: &Id) {
309 if std::thread::panicking() {
310 // `exit()` can be called in `drop` impls, so we must guard against
311 // double panics.
312 println!("[{}] exit {:?} while panicking", self.name, id);
313 return;
314 }
315 let spans = self.spans.lock().unwrap();
316 let span = spans
317 .get(id)
318 .unwrap_or_else(|| panic!("[{}] no span for ID {:?}", self.name, id));
319 println!("[{}] exit: {}; id={:?};", self.name, span.name, id);
320 match self.expected.lock().unwrap().pop_front() {
321 None => {}
322 Some(Expect::Exit(ref expected_span)) => {
323 if let Some(name) = expected_span.name() {
324 assert_eq!(name, span.name);
325 }
326 let curr = self.current.lock().unwrap().pop();
327 assert_eq!(
328 Some(id),
329 curr.as_ref(),
330 "[{}] exited span {:?}, but the current span was {:?}",
331 self.name,
332 span.name,
333 curr.as_ref().and_then(|id| spans.get(id)).map(|s| s.name)
334 );
335 }
336 Some(ex) => ex.bad(&self.name, format_args!("exited span {:?}", span.name)),
337 };
338 }
339
clone_span(&self, id: &Id) -> Id340 fn clone_span(&self, id: &Id) -> Id {
341 let name = self.spans.lock().unwrap().get_mut(id).map(|span| {
342 let name = span.name;
343 println!(
344 "[{}] clone_span: {}; id={:?}; refs={:?};",
345 self.name, name, id, span.refs
346 );
347 span.refs += 1;
348 name
349 });
350 if name.is_none() {
351 println!("[{}] clone_span: id={:?};", self.name, id);
352 }
353 let mut expected = self.expected.lock().unwrap();
354 let was_expected = if let Some(Expect::CloneSpan(ref span)) = expected.front() {
355 assert_eq!(
356 name,
357 span.name(),
358 "[{}] expected to clone a span named {:?}",
359 self.name,
360 span.name()
361 );
362 true
363 } else {
364 false
365 };
366 if was_expected {
367 expected.pop_front();
368 }
369 id.clone()
370 }
371
drop_span(&self, id: Id)372 fn drop_span(&self, id: Id) {
373 let mut is_event = false;
374 let name = if let Ok(mut spans) = self.spans.try_lock() {
375 spans.get_mut(&id).map(|span| {
376 let name = span.name;
377 if name.contains("event") {
378 is_event = true;
379 }
380 println!(
381 "[{}] drop_span: {}; id={:?}; refs={:?};",
382 self.name, name, id, span.refs
383 );
384 span.refs -= 1;
385 name
386 })
387 } else {
388 None
389 };
390 if name.is_none() {
391 println!("[{}] drop_span: id={:?}", self.name, id);
392 }
393 if let Ok(mut expected) = self.expected.try_lock() {
394 let was_expected = match expected.front() {
395 Some(Expect::DropSpan(ref span)) => {
396 // Don't assert if this function was called while panicking,
397 // as failing the assertion can cause a double panic.
398 if !::std::thread::panicking() {
399 assert_eq!(name, span.name());
400 }
401 true
402 }
403 Some(Expect::Event(_)) => {
404 if !::std::thread::panicking() {
405 assert!(is_event, "[{}] expected an event", self.name);
406 }
407 true
408 }
409 _ => false,
410 };
411 if was_expected {
412 expected.pop_front();
413 }
414 }
415 }
416 }
417
418 impl MockHandle {
new(expected: Arc<Mutex<VecDeque<Expect>>>, name: String) -> Self419 pub fn new(expected: Arc<Mutex<VecDeque<Expect>>>, name: String) -> Self {
420 Self(expected, name)
421 }
422
assert_finished(&self)423 pub fn assert_finished(&self) {
424 if let Ok(ref expected) = self.0.lock() {
425 assert!(
426 !expected.iter().any(|thing| thing != &Expect::Nothing),
427 "\n[{}] more notifications expected: {:#?}",
428 self.1,
429 **expected
430 );
431 }
432 }
433 }
434
435 impl Expect {
bad(&self, name: impl AsRef<str>, what: fmt::Arguments<'_>)436 pub fn bad(&self, name: impl AsRef<str>, what: fmt::Arguments<'_>) {
437 let name = name.as_ref();
438 match self {
439 Expect::Event(e) => panic!(
440 "\n[{}] expected event {}\n[{}] but instead {}",
441 name, e, name, what,
442 ),
443 Expect::Enter(e) => panic!(
444 "\n[{}] expected to enter {}\n[{}] but instead {}",
445 name, e, name, what,
446 ),
447 Expect::Exit(e) => panic!(
448 "\n[{}] expected to exit {}\n[{}] but instead {}",
449 name, e, name, what,
450 ),
451 Expect::CloneSpan(e) => {
452 panic!(
453 "\n[{}] expected to clone {}\n[{}] but instead {}",
454 name, e, name, what,
455 )
456 }
457 Expect::DropSpan(e) => {
458 panic!(
459 "\n[{}] expected to drop {}\n[{}] but instead {}",
460 name, e, name, what,
461 )
462 }
463 Expect::Visit(e, fields) => panic!(
464 "\n[{}] expected {} to record {}\n[{}] but instead {}",
465 name, e, fields, name, what,
466 ),
467 Expect::NewSpan(e) => panic!(
468 "\n[{}] expected {}\n[{}] but instead {}",
469 name, e, name, what
470 ),
471 Expect::Nothing => panic!(
472 "\n[{}] expected nothing else to happen\n[{}] but {} instead",
473 name, name, what,
474 ),
475 }
476 }
477 }
478