mock: match parent span on ExpectedSpan (#3098)

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.

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/event.rs

@@ -29,7 +29,12 @@
 //! [`subscriber`]: mod@crate::subscriber
 //! [`expect::event`]: fn@crate::expect::event
 #![allow(missing_docs)]
-use crate::{ancestry::Ancestry, expect, field, metadata::ExpectedMetadata, span};
+use crate::{
+    ancestry::{ActualAncestry, ExpectedAncestry},
+    expect, field,
+    metadata::ExpectedMetadata,
+    span,
+};
 
 use std::fmt;
 
@@ -42,7 +47,7 @@ use std::fmt;
 #[derive(Default, Eq, PartialEq)]
 pub struct ExpectedEvent {
     pub(super) fields: Option<field::ExpectedFields>,
-    pub(super) ancestry: Option<Ancestry>,
+    pub(super) ancestry: Option<ExpectedAncestry>,
     pub(super) in_spans: Option<Vec<span::ExpectedSpan>>,
     pub(super) metadata: ExpectedMetadata,
 }
@@ -253,9 +258,10 @@ impl ExpectedEvent {
         }
     }
 
-    /// Configures this `ExpectedEvent` to expect the specified [`Ancestry`].
-    /// An event's ancestry indicates whether is has a parent or is a root, and
-    /// whether the parent is explicitly or contextually assigned.
+    /// Configures this `ExpectedEvent` to expect the specified
+    /// [`ExpectedAncestry`]. An event's ancestry indicates whether is has a
+    /// parent or is a root, and whether the parent is explicitly or
+    /// contextually assigned.
     ///
     /// An _explicit_ parent span is one passed to the `event!` macro in the
     /// `parent:` field. If no `parent:` field is specified, then the event
@@ -267,9 +273,34 @@ impl ExpectedEvent {
     ///
     /// # Examples
     ///
-    /// If `expect::has_explicit_parent("parent_name")` is passed
-    /// `with_ancestry` then the provided string is the name of the explicit
-    /// parent span to expect.
+    /// An explicit or contextual can be matched on an `ExpectedSpan`.
+    ///
+    /// ```
+    /// use tracing::subscriber::with_default;
+    /// use tracing_mock::{subscriber, expect};
+    ///
+    /// let parent = expect::span()
+    ///     .named("parent_span")
+    ///     .with_target("custom-target")
+    ///     .at_level(tracing::Level::INFO);
+    /// let event = expect::event()
+    ///     .with_ancestry(expect::has_explicit_parent(parent));
+    ///
+    /// let (subscriber, handle) = subscriber::mock()
+    ///     .event(event)
+    ///     .run_with_handle();
+    ///
+    /// with_default(subscriber, || {
+    ///     let parent = tracing::info_span!(target: "custom-target", "parent_span");
+    ///     tracing::info!(parent: parent.id(), field = &"value");
+    /// });
+    ///
+    /// handle.assert_finished();
+    /// ```
+    /// The functions `expect::has_explicit_parent` and
+    /// `expect::has_contextual_parent` take `Into<ExpectedSpan>`, so a string
+    /// passed directly will match on a span with that name, or an
+    /// [`ExpectedId`] can be passed to match a span with that Id.
     ///
     /// ```
     /// use tracing::subscriber::with_default;
@@ -382,7 +413,9 @@ impl ExpectedEvent {
     ///
     /// handle.assert_finished();
     /// ```
-    pub fn with_ancestry(self, ancenstry: Ancestry) -> ExpectedEvent {
+    ///
+    /// [`ExpectedId`]: struct@crate::span::ExpectedId
+    pub fn with_ancestry(self, ancenstry: ExpectedAncestry) -> ExpectedEvent {
         Self {
             ancestry: Some(ancenstry),
             ..self
@@ -506,7 +539,7 @@ impl ExpectedEvent {
     pub(crate) fn check(
         &mut self,
         event: &tracing::Event<'_>,
-        get_ancestry: impl FnOnce() -> Ancestry,
+        get_ancestry: impl FnOnce() -> ActualAncestry,
         subscriber_name: &str,
     ) {
         let meta = event.metadata();

tracing-mock/src/expect.rs

@@ -1,7 +1,7 @@
 use std::fmt;
 
 use crate::{
-    ancestry::Ancestry,
+    ancestry::ExpectedAncestry,
     event::ExpectedEvent,
     field::{ExpectedField, ExpectedFields, ExpectedValue},
     span::{ExpectedId, ExpectedSpan, NewSpan},
@@ -71,26 +71,26 @@ pub fn id() -> ExpectedId {
     ExpectedId::new_unset()
 }
 
-/// Convenience function that returns [`Ancestry::IsContextualRoot`].
-pub fn is_contextual_root() -> Ancestry {
-    Ancestry::IsContextualRoot
+/// Convenience function that returns [`ExpectedAncestry::IsContextualRoot`].
+pub fn is_contextual_root() -> ExpectedAncestry {
+    ExpectedAncestry::IsContextualRoot
 }
 
-/// Convenience function that returns [`Ancestry::HasContextualParent`] with
+/// Convenience function that returns [`ExpectedAncestry::HasContextualParent`] with
 /// provided name.
-pub fn has_contextual_parent<S: Into<String>>(name: S) -> Ancestry {
-    Ancestry::HasContextualParent(name.into())
+pub fn has_contextual_parent<S: Into<ExpectedSpan>>(span: S) -> ExpectedAncestry {
+    ExpectedAncestry::HasContextualParent(span.into())
 }
 
-/// Convenience function that returns [`Ancestry::IsExplicitRoot`].
-pub fn is_explicit_root() -> Ancestry {
-    Ancestry::IsExplicitRoot
+/// Convenience function that returns [`ExpectedAncestry::IsExplicitRoot`].
+pub fn is_explicit_root() -> ExpectedAncestry {
+    ExpectedAncestry::IsExplicitRoot
 }
 
-/// Convenience function that returns [`Ancestry::HasExplicitParent`] with
+/// Convenience function that returns [`ExpectedAncestry::HasExplicitParent`] with
 /// provided name.
-pub fn has_explicit_parent<S: Into<String>>(name: S) -> Ancestry {
-    Ancestry::HasExplicitParent(name.into())
+pub fn has_explicit_parent<S: Into<ExpectedSpan>>(span: S) -> ExpectedAncestry {
+    ExpectedAncestry::HasExplicitParent(span.into())
 }
 
 impl Expect {