mock: correct contextual/explicit parent assertions

tracing-attributes/tests/parents.rs

@@ -21,23 +21,16 @@ fn default_parent_test() {
         .new_span(
             contextual_parent
                 .clone()
-                .with_contextual_parent(None)
-                .with_explicit_parent(None),
-        )
-        .new_span(
-            child
-                .clone()
-                .with_contextual_parent(Some("contextual_parent"))
-                .with_explicit_parent(None),
+                .with_ancestry(expect::is_contextual_root()),
         )
+        .new_span(child.clone().with_ancestry(expect::is_contextual_root()))
         .enter(child.clone())
         .exit(child.clone())
         .enter(contextual_parent.clone())
         .new_span(
             child
                 .clone()
-                .with_contextual_parent(Some("contextual_parent"))
-                .with_explicit_parent(None),
+                .with_ancestry(expect::has_contextual_parent("contextual_parent")),
         )
         .enter(child.clone())
         .exit(child)
@@ -68,20 +61,14 @@ fn explicit_parent_test() {
         .new_span(
             contextual_parent
                 .clone()
-                .with_contextual_parent(None)
-                .with_explicit_parent(None),
-        )
-        .new_span(
-            explicit_parent
-                .with_contextual_parent(None)
-                .with_explicit_parent(None),
+                .with_ancestry(expect::is_contextual_root()),
         )
+        .new_span(explicit_parent.with_ancestry(expect::is_contextual_root()))
         .enter(contextual_parent.clone())
         .new_span(
             child
                 .clone()
-                .with_contextual_parent(Some("contextual_parent"))
-                .with_explicit_parent(Some("explicit_parent")),
+                .with_ancestry(expect::has_explicit_parent("explicit_parent")),
         )
         .enter(child.clone())
         .exit(child)

tracing-futures/tests/std_future.rs

@@ -71,7 +71,7 @@ fn span_on_drop() {
         .enter(expect::span().named("foo"))
         .event(
             expect::event()
-                .with_contextual_parent(Some("foo"))
+                .with_ancestry(expect::has_contextual_parent("foo"))
                 .at_level(Level::INFO),
         )
         .exit(expect::span().named("foo"))
@@ -81,7 +81,7 @@ fn span_on_drop() {
         .enter(expect::span().named("bar"))
         .event(
             expect::event()
-                .with_contextual_parent(Some("bar"))
+                .with_ancestry(expect::has_contextual_parent("bar"))
                 .at_level(Level::INFO),
         )
         .exit(expect::span().named("bar"))

tracing-mock/src/ancestry.rs

@@ -0,0 +1,145 @@
+//! Define the ancestry of an event or span.
+//!
+//! See the documentation on the [`Ancestry`] enum for further details.
+
+use tracing_core::{
+    span::{self, Attributes},
+    Event,
+};
+
+/// 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),
+    /// 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. Additionally,
+    /// it has no explicitly assigned parent.
+    HasContextualParent(String),
+    /// The event or span is a contextual root. It has no contextual parent and also has no
+    /// explicitly assigned parent.
+    IsContextualRoot,
+}
+
+impl Ancestry {
+    #[track_caller]
+    pub(crate) fn check(
+        &self,
+        actual_ancestry: &Ancestry,
+        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}'")
+            }
+        };
+
+        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::IsContextualRoot => "was actually a contextual root".to_string(),
+            Self::HasContextualParent(name) => {
+                format!("actually has a contextual parent with name='{name}'")
+            }
+        };
+
+        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)
+        );
+    }
+}
+
+pub(crate) trait HasAncestry {
+    fn is_contextual(&self) -> bool;
+
+    fn is_root(&self) -> bool;
+
+    fn parent(&self) -> Option<&span::Id>;
+}
+
+impl HasAncestry for &Event<'_> {
+    fn is_contextual(&self) -> bool {
+        (self as &Event<'_>).is_contextual()
+    }
+
+    fn is_root(&self) -> bool {
+        (self as &Event<'_>).is_root()
+    }
+
+    fn parent(&self) -> Option<&span::Id> {
+        (self as &Event<'_>).parent()
+    }
+}
+
+impl HasAncestry for &Attributes<'_> {
+    fn is_contextual(&self) -> bool {
+        (self as &Attributes<'_>).is_contextual()
+    }
+
+    fn is_root(&self) -> bool {
+        (self as &Attributes<'_>).is_root()
+    }
+
+    fn parent(&self) -> Option<&span::Id> {
+        (self as &Attributes<'_>).parent()
+    }
+}
+
+/// Determines the ancestry of an actual span or event.
+///
+/// The rules for determining the ancestry are as follows:
+///
+/// +------------+--------------+-----------------+---------------------+
+/// | Contextual | Current Span | Explicit Parent | Ancestry            |
+/// +------------+--------------+-----------------+---------------------+
+/// | Yes        | Yes          | -               | HasContextualParent |
+/// | Yes        | No           | -               | IsContextualRoot    |
+/// | No         | -            | Yes             | HasExplicitParent   |
+/// | No         | -            | No              | IsExplicitRoot      |
+/// +------------+--------------+-----------------+---------------------+
+pub(crate) fn get_ancestry(
+    item: impl HasAncestry,
+    lookup_current: impl FnOnce() -> Option<span::Id>,
+    span_name: impl FnOnce(&span::Id) -> Option<&str>,
+) -> Ancestry {
+    if item.is_contextual() {
+        if let Some(parent_id) = lookup_current() {
+            let contextual_parent_name = span_name(&parent_id).expect(
+                "tracing-mock: contextual parent cannot \
+                            be looked up by ID. Was it recorded correctly?",
+            );
+            Ancestry::HasContextualParent(contextual_parent_name.to_string())
+        } else {
+            Ancestry::IsContextualRoot
+        }
+    } else if item.is_root() {
+        Ancestry::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(
+            "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())
+    }
+}

tracing-mock/src/collector.rs

@@ -138,6 +138,7 @@
 //! [`Collect`]: trait@tracing::Collect
 //! [`MockCollector`]: struct@crate::collector::MockCollector
 use crate::{
+    ancestry::get_ancestry,
     event::ExpectedEvent,
     expect::Expect,
     field::ExpectedFields,
@@ -1039,16 +1040,20 @@ where
                         )
                     }
                 }
-                let get_parent_name = || {
-                    let stack = self.current.lock().unwrap();
-                    let spans = self.spans.lock().unwrap();
-                    event
-                        .parent()
-                        .and_then(|id| spans.get(id))
-                        .or_else(|| stack.last().and_then(|id| spans.get(id)))
-                        .map(|s| s.name.to_string())
+                let event_get_ancestry = || {
+                    get_ancestry(
+                        event,
+                        || self.lookup_current(),
+                        |span_id| {
+                            self.spans
+                                .lock()
+                                .unwrap()
+                                .get(span_id)
+                                .map(|span| span.name)
+                        },
+                    )
                 };
-                expected.check(event, get_parent_name, &self.name);
+                expected.check(event, event_get_ancestry, &self.name);
             }
             Some(ex) => ex.bad(&self.name, format_args!("observed event {:#?}", event)),
         }
@@ -1108,14 +1113,18 @@ where
                 if let Some(expected_id) = &expected.span.id {
                     expected_id.set(id.into_u64()).unwrap();
                 }
-                let get_parent_name = || {
-                    let stack = self.current.lock().unwrap();
-                    span.parent()
-                        .and_then(|id| spans.get(id))
-                        .or_else(|| stack.last().and_then(|id| spans.get(id)))
-                        .map(|s| s.name.to_string())
-                };
-                expected.check(span, get_parent_name, &self.name);
+
+                expected.check(
+                    span,
+                    || {
+                        get_ancestry(
+                            span,
+                            || self.lookup_current(),
+                            |span_id| spans.get(span_id).map(|span| span.name),
+                        )
+                    },
+                    &self.name,
+                );
             }
         }
         spans.insert(
@@ -1265,6 +1274,16 @@ where
     }
 }
 
+impl<F> Running<F>
+where
+    F: Fn(&Metadata<'_>) -> bool,
+{
+    fn lookup_current(&self) -> Option<span::Id> {
+        let stack = self.current.lock().unwrap();
+        stack.last().cloned()
+    }
+}
+
 impl MockHandle {
     #[cfg(feature = "tracing-subscriber")]
     pub(crate) fn new(expected: Arc<Mutex<VecDeque<Expect>>>, name: String) -> Self {