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