mock: match parent span on ExpectedSpan (#3098)

## Motivation

The `with_ancestry` methods on `NewSpan` and `ExpectedEvent` provide a
way to match whether the span or event is a contextual or explicit root
or if it has a contextual or explicit parent span.

However, in the case of matching on a contextual or explicit parent
span, only the span name could be used for matching. This is
sufficiently precise when testing tracing instrumentation in other
libraries or applications as opposed to testing tracing itself.

It is likely that a user would like to test that some span or event has
a specific span as a parent, and not just any span with a specific name,
in many cases, all the possible parent spans may have the same name.
This is the case when testing tracing instrumentation in Tokio.

## Solution

To solve this problem, the `Ancestry` struct was renamed to
`ExpectedAncestry` and in the case of expecting an explicit or
conextual parent, an `ExpectedSpan` object can be passed in. This
provides the maximum possible flexibility.

The convenience functions in the `expect` module now take
`Into<ExpectedSpan>` so that existing tests that pass a string type
object for the parent will see the same behaviour as previously and
shorthand use for expected Ids is also available.

Additionally, the span checking code has been unified between the
`MockCollector` and `MockSubscriber` cases and the assertion
descriptions have been improved to make them more readable.

tracing-mock/src/ancestry.rs

@@ -1,71 +1,102 @@
 //! Define the ancestry of an event or span.
 //!
-//! See the documentation on the [`Ancestry`] enum for further details.
+//! See the documentation on the [`ExpectedAncestry`] enum for further details.
 
 use tracing_core::{
     span::{self, Attributes},
     Event,
 };
 
+use crate::span::{ActualSpan, ExpectedSpan};
+
 /// The ancestry of an event or span.
 ///
 /// An event or span can have an explicitly assigned parent, or be an explicit root. Otherwise,
 /// an event or span may have a contextually assigned parent or in the final case will be a
 /// contextual root.
 #[derive(Debug, Eq, PartialEq)]
-pub enum Ancestry {
-    /// The event or span has an explicitly assigned parent (created with `parent: span_id`) with
-    /// the specified name.
-    HasExplicitParent(String),
+pub enum ExpectedAncestry {
+    /// The event or span has an explicitly assigned parent (created with `parent: span_id`) span.
+    HasExplicitParent(ExpectedSpan),
     /// The event or span is an explicitly defined root. It was created with `parent: None` and
     /// has no parent.
     IsExplicitRoot,
-    /// The event or span has a contextually assigned parent with the specified name. It has no
-    /// explicitly assigned parent, nor has it been explicitly defined as a root (it was created
-    /// without the `parent:` directive). There was a span in context when this event or span was
-    /// created.
-    HasContextualParent(String),
+    /// The event or span has a contextually assigned parent span. It has no explicitly assigned
+    /// parent span, nor has it been explicitly defined as a root (it was created without the
+    /// `parent:` directive). There was a span in context when this event or span was created.
+    HasContextualParent(ExpectedSpan),
     /// The event or span is a contextual root. It has no explicitly assigned parent, nor has it
     /// been explicitly defined as a root (it was created without the `parent:` directive).
     /// Additionally, no span was in context when this event or span was created.
     IsContextualRoot,
 }
 
-impl Ancestry {
+pub(crate) enum ActualAncestry {
+    HasExplicitParent(ActualSpan),
+    IsExplicitRoot,
+    HasContextualParent(ActualSpan),
+    IsContextualRoot,
+}
+
+impl ExpectedAncestry {
     #[track_caller]
     pub(crate) fn check(
         &self,
-        actual_ancestry: &Ancestry,
+        actual_ancestry: &ActualAncestry,
         ctx: impl std::fmt::Display,
         collector_name: &str,
     ) {
-        let expected_description = |ancestry: &Ancestry| match ancestry {
-            Self::IsExplicitRoot => "be an explicit root".to_string(),
-            Self::HasExplicitParent(name) => format!("have an explicit parent with name='{name}'"),
-            Self::IsContextualRoot => "be a contextual root".to_string(),
-            Self::HasContextualParent(name) => {
-                format!("have a contextual parent with name='{name}'")
+        match (self, actual_ancestry) {
+            (Self::IsExplicitRoot, ActualAncestry::IsExplicitRoot) => {}
+            (Self::IsContextualRoot, ActualAncestry::IsContextualRoot) => {}
+            (
+                Self::HasExplicitParent(expected_parent),
+                ActualAncestry::HasExplicitParent(actual_parent),
+            ) => {
+                expected_parent.check(
+                    actual_parent,
+                    format_args!("{ctx} to have an explicit parent span"),
+                    collector_name,
+                );
             }
-        };
-
-        let actual_description = |ancestry: &Ancestry| match ancestry {
-            Self::IsExplicitRoot => "was actually an explicit root".to_string(),
-            Self::HasExplicitParent(name) => {
-                format!("actually has an explicit parent with name='{name}'")
+            (
+                Self::HasContextualParent(expected_parent),
+                ActualAncestry::HasContextualParent(actual_parent),
+            ) => {
+                println!("----> [{collector_name}] check {expected_parent:?} against actual parent with Id={id:?}", id = actual_parent.id());
+                expected_parent.check(
+                    actual_parent,
+                    format_args!("{ctx} to have a contextual parent span"),
+                    collector_name,
+                );
             }
-            Self::IsContextualRoot => "was actually a contextual root".to_string(),
-            Self::HasContextualParent(name) => {
-                format!("actually has a contextual parent with name='{name}'")
+            _ => {
+                // Ancestry types don't match at all.
+                let expected_description = match self {
+                    Self::IsExplicitRoot => "be an explicit root",
+                    Self::HasExplicitParent(_) => "have an explicit parent span",
+                    Self::IsContextualRoot => "be a contextual root",
+                    Self::HasContextualParent(_) => "have a contextual parent span",
+                };
+
+                let actual_description = match actual_ancestry {
+                    ActualAncestry::IsExplicitRoot => "is actually an explicit root",
+                    ActualAncestry::HasExplicitParent(_) => "actually has an explicit parent span",
+                    ActualAncestry::IsContextualRoot => "is actually a contextual root",
+                    ActualAncestry::HasContextualParent(_) => {
+                        "actually has a contextual parent span"
+                    }
+                };
+
+                panic!(
+                    "{}",
+                    format!(
+                        "[{collector_name}] expected {ctx} to {expected_description}, \
+                        but it {actual_description}"
+                    )
+                );
             }
-        };
-
-        assert_eq!(
-            self,
-            actual_ancestry,
-            "[{collector_name}] expected {ctx} to {expected_description}, but {actual_description}",
-            expected_description = expected_description(self),
-            actual_description = actual_description(actual_ancestry)
-        );
+        }
     }
 }
 
@@ -120,29 +151,29 @@ impl HasAncestry for &Attributes<'_> {
 pub(crate) fn get_ancestry(
     item: impl HasAncestry,
     lookup_current: impl FnOnce() -> Option<span::Id>,
-    span_name: impl FnOnce(&span::Id) -> Option<&str>,
-) -> Ancestry {
+    actual_span: impl FnOnce(&span::Id) -> Option<ActualSpan>,
+) -> ActualAncestry {
     if item.is_contextual() {
         if let Some(parent_id) = lookup_current() {
-            let contextual_parent_name = span_name(&parent_id).expect(
+            let contextual_parent_span = actual_span(&parent_id).expect(
                 "tracing-mock: contextual parent cannot \
                             be looked up by ID. Was it recorded correctly?",
             );
-            Ancestry::HasContextualParent(contextual_parent_name.to_string())
+            ActualAncestry::HasContextualParent(contextual_parent_span)
         } else {
-            Ancestry::IsContextualRoot
+            ActualAncestry::IsContextualRoot
         }
     } else if item.is_root() {
-        Ancestry::IsExplicitRoot
+        ActualAncestry::IsExplicitRoot
     } else {
         let parent_id = item.parent().expect(
             "tracing-mock: is_contextual=false is_root=false \
                         but no explicit parent found. This is a bug!",
         );
-        let explicit_parent_name = span_name(parent_id).expect(
+        let explicit_parent_span = actual_span(parent_id).expect(
             "tracing-mock: explicit parent cannot be looked \
                         up by ID. Is the provided Span ID valid: {parent_id}",
         );
-        Ancestry::HasExplicitParent(explicit_parent_name.to_string())
+        ActualAncestry::HasExplicitParent(explicit_parent_span)
     }
 }

tracing-mock/src/collector.rs

@@ -142,7 +142,7 @@ use crate::{
     event::ExpectedEvent,
     expect::Expect,
     field::ExpectedFields,
-    span::{ExpectedSpan, NewSpan},
+    span::{ActualSpan, ExpectedSpan, NewSpan},
 };
 use std::{
     collections::{HashMap, VecDeque},
@@ -160,19 +160,15 @@ use tracing::{
 };
 
 pub(crate) struct SpanState {
-    id: u64,
+    id: Id,
     name: &'static str,
     refs: usize,
     meta: &'static Metadata<'static>,
 }
 
-impl SpanState {
-    pub(crate) fn id(&self) -> u64 {
-        self.id
-    }
-
-    pub(crate) fn metadata(&self) -> &'static Metadata<'static> {
-        self.meta
+impl From<&SpanState> for ActualSpan {
+    fn from(span_state: &SpanState) -> Self {
+        Self::new(span_state.id.clone(), Some(span_state.meta))
     }
 }
 
@@ -1069,7 +1065,7 @@ where
                                 .lock()
                                 .unwrap()
                                 .get(span_id)
-                                .map(|span| span.name)
+                                .map(|span| span.into())
                         },
                     )
                 };
@@ -1140,7 +1136,7 @@ where
                         get_ancestry(
                             span,
                             || self.lookup_current(),
-                            |span_id| spans.get(span_id).map(|span| span.name),
+                            |span_id| spans.get(span_id).map(|span| span.into()),
                         )
                     },
                     &self.name,
@@ -1150,7 +1146,7 @@ where
         spans.insert(
             id.clone(),
             SpanState {
-                id: id.into_u64(),
+                id: id.clone(),
                 name: meta.name(),
                 refs: 1,
                 meta,
@@ -1166,7 +1162,7 @@ where
             match self.expected.lock().unwrap().pop_front() {
                 None => {}
                 Some(Expect::Enter(ref expected_span)) => {
-                    expected_span.check(span, &self.name);
+                    expected_span.check(&span.into(), "to enter a span", &self.name);
                 }
                 Some(ex) => ex.bad(&self.name, format_args!("entered span {:?}", span.name)),
             }
@@ -1189,7 +1185,7 @@ where
         match self.expected.lock().unwrap().pop_front() {
             None => {}
             Some(Expect::Exit(ref expected_span)) => {
-                expected_span.check(span, &self.name);
+                expected_span.check(&span.into(), "to exit a span", &self.name);
                 let curr = self.current.lock().unwrap().pop();
                 assert_eq!(
                     Some(id),
@@ -1205,27 +1201,34 @@ where
     }
 
     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={:?};",
-                self.name, name, id, span.refs
-            );
-            span.refs += 1;
-            name
-        });
-        if name.is_none() {
-            println!("[{}] clone_span: id={:?};", self.name, 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 span)) = expected.front() {
-            assert_eq!(
-                name,
-                span.name(),
-                "[{}] expected to clone a span named {:?}",
-                self.name,
-                span.name()
-            );
+        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);
+                }
+                // Check only by Id
+                None => expected_span.check(&id.into(), "to clone a span", &self.name),
+            }
             true
         } else {
             false