use std::{
error, fmt,
sync::{
atomic::{AtomicU64, Ordering},
Arc,
},
};
use crate::{
ancestry::{ActualAncestry, ExpectedAncestry},
field::ExpectedFields,
metadata::ExpectedMetadata,
};
#[derive(Clone, Default, Eq, PartialEq)]
pub struct ExpectedSpan {
pub(crate) id: Option<ExpectedId>,
pub(crate) metadata: ExpectedMetadata,
}
impl<I> From<I> for ExpectedSpan
where
I: Into<String>,
{
fn from(name: I) -> Self {
ExpectedSpan::default().named(name)
}
}
impl From<&ExpectedId> for ExpectedSpan {
fn from(id: &ExpectedId) -> Self {
ExpectedSpan::default().with_id(id.clone())
}
}
impl From<&ExpectedSpan> for ExpectedSpan {
fn from(span: &ExpectedSpan) -> Self {
span.clone()
}
}
#[derive(Default, Eq, PartialEq)]
pub struct NewSpan {
pub(crate) span: ExpectedSpan,
pub(crate) fields: ExpectedFields,
pub(crate) ancestry: Option<ExpectedAncestry>,
}
pub(crate) struct ActualSpan {
id: tracing_core::span::Id,
metadata: Option<&'static tracing_core::Metadata<'static>>,
}
impl ActualSpan {
pub(crate) fn new(
id: tracing_core::span::Id,
metadata: Option<&'static tracing_core::Metadata<'static>>,
) -> Self {
Self { id, metadata }
}
pub(crate) fn id(&self) -> tracing_core::span::Id {
self.id.clone()
}
pub(crate) fn metadata(&self) -> Option<&'static tracing_core::Metadata<'static>> {
self.metadata
}
}
impl From<&tracing_core::span::Id> for ActualSpan {
fn from(id: &tracing_core::span::Id) -> Self {
Self::new(id.clone(), None)
}
}
#[derive(Clone, Default)]
pub struct ExpectedId {
inner: Arc<AtomicU64>,
}
impl ExpectedSpan {
pub fn named<I>(self, name: I) -> Self
where
I: Into<String>,
{
Self {
metadata: ExpectedMetadata {
name: Some(name.into()),
..self.metadata
},
..self
}
}
pub fn with_id(self, id: ExpectedId) -> Self {
Self {
id: Some(id),
..self
}
}
pub fn at_level(self, level: tracing::Level) -> Self {
Self {
metadata: ExpectedMetadata {
level: Some(level),
..self.metadata
},
..self
}
}
pub fn with_target<I>(self, target: I) -> Self
where
I: Into<String>,
{
Self {
metadata: ExpectedMetadata {
target: Some(target.into()),
..self.metadata
},
..self
}
}
pub fn with_ancestry(self, ancestry: ExpectedAncestry) -> NewSpan {
NewSpan {
ancestry: Some(ancestry),
span: self,
..Default::default()
}
}
pub fn with_fields<I>(self, fields: I) -> NewSpan
where
I: Into<ExpectedFields>,
{
NewSpan {
span: self,
fields: fields.into(),
..Default::default()
}
}
pub(crate) fn id(&self) -> Option<&ExpectedId> {
self.id.as_ref()
}
pub(crate) fn name(&self) -> Option<&str> {
self.metadata.name.as_ref().map(String::as_ref)
}
pub(crate) fn level(&self) -> Option<tracing::Level> {
self.metadata.level
}
pub(crate) fn target(&self) -> Option<&str> {
self.metadata.target.as_deref()
}
pub(crate) fn check(&self, actual: &ActualSpan, ctx: impl fmt::Display, subscriber_name: &str) {
if let Some(expected_id) = &self.id {
expected_id.check(&actual.id(), format_args!("{ctx} a span"), subscriber_name);
}
match actual.metadata() {
Some(actual_metadata) => self.metadata.check(actual_metadata, ctx, subscriber_name),
None => {
if self.metadata.has_expectations() {
panic!(
"{}",
format_args!(
"[{subscriber_name}] expected {ctx} a span with valid metadata, \
but got one with unknown Id={actual_id}",
actual_id = actual.id().into_u64()
)
);
}
}
}
}
}
impl fmt::Debug for ExpectedSpan {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut s = f.debug_struct("MockSpan");
if let Some(id) = self.id() {
s.field("id", &id);
}
if let Some(name) = self.name() {
s.field("name", &name);
}
if let Some(level) = self.level() {
s.field("level", &format_args!("{:?}", level));
}
if let Some(target) = self.target() {
s.field("target", &target);
}
s.finish()
}
}
impl fmt::Display for ExpectedSpan {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.metadata.name.is_some() {
write!(f, "a span{}", self.metadata)
} else {
write!(f, "any span{}", self.metadata)
}
}
}
impl<S> From<S> for NewSpan
where
S: Into<ExpectedSpan>,
{
fn from(span: S) -> Self {
Self {
span: span.into(),
..Default::default()
}
}
}
impl NewSpan {
pub fn with_ancestry(self, ancestry: ExpectedAncestry) -> NewSpan {
NewSpan {
ancestry: Some(ancestry),
..self
}
}
pub fn with_fields<I>(self, fields: I) -> NewSpan
where
I: Into<ExpectedFields>,
{
NewSpan {
fields: fields.into(),
..self
}
}
pub(crate) fn check(
&mut self,
span: &tracing_core::span::Attributes<'_>,
get_ancestry: impl FnOnce() -> ActualAncestry,
subscriber_name: &str,
) {
let meta = span.metadata();
let name = meta.name();
self.span
.metadata
.check(meta, "a new span", subscriber_name);
let mut checker = self.fields.checker(name, subscriber_name);
span.record(&mut checker);
checker.finish();
if let Some(ref expected_ancestry) = self.ancestry {
let actual_ancestry = get_ancestry();
expected_ancestry.check(
&actual_ancestry,
format_args!("span `{}`", name),
subscriber_name,
);
}
}
}
impl fmt::Display for NewSpan {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "a new span{}", self.span.metadata)?;
if !self.fields.is_empty() {
write!(f, " with {}", self.fields)?;
}
Ok(())
}
}
impl fmt::Debug for NewSpan {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut s = f.debug_struct("NewSpan");
if let Some(name) = self.span.name() {
s.field("name", &name);
}
if let Some(level) = self.span.level() {
s.field("level", &format_args!("{:?}", level));
}
if let Some(target) = self.span.target() {
s.field("target", &target);
}
if let Some(ref parent) = self.ancestry {
s.field("parent", &format_args!("{:?}", parent));
}
if !self.fields.is_empty() {
s.field("fields", &self.fields);
}
s.finish()
}
}
impl PartialEq for ExpectedId {
fn eq(&self, other: &Self) -> bool {
self.inner.load(Ordering::Relaxed) == other.inner.load(Ordering::Relaxed)
}
}
impl Eq for ExpectedId {}
impl fmt::Debug for ExpectedId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("ExpectedId").field(&self.inner).finish()
}
}
impl ExpectedId {
const UNSET: u64 = 0;
pub(crate) fn new_unset() -> Self {
Self {
inner: Arc::new(AtomicU64::from(Self::UNSET)),
}
}
pub(crate) fn set(&self, span_id: u64) -> Result<(), SetActualSpanIdError> {
self.inner
.compare_exchange(Self::UNSET, span_id, Ordering::Relaxed, Ordering::Relaxed)
.map_err(|current| SetActualSpanIdError {
previous_span_id: current,
new_span_id: span_id,
})?;
Ok(())
}
pub(crate) fn check(
&self,
actual: &tracing_core::span::Id,
ctx: fmt::Arguments<'_>,
subscriber_name: &str,
) {
let expected_id = self.inner.load(Ordering::Relaxed);
let actual_id = actual.into_u64();
assert!(
expected_id != Self::UNSET,
"{}",
format!(
"\n[{subscriber_name}] expected {ctx} with an expected Id set,\n\
[{subscriber_name}] but it hasn't been, perhaps this `ExpectedId` \
wasn't used in a call to `new_span()`?"
)
);
assert_eq!(
expected_id,
actual_id,
"{}",
format_args!(
"\n[{subscriber_name}] expected {ctx} with Id `{expected_id}`,\n\
[{subscriber_name}] but got one with Id `{actual_id}` instead",
)
);
}
}
#[derive(Debug)]
pub(crate) struct SetActualSpanIdError {
previous_span_id: u64,
new_span_id: u64,
}
impl fmt::Display for SetActualSpanIdError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"Could not set `ExpecedId` to {new}, \
it had already been set to {previous}",
new = self.new_span_id,
previous = self.previous_span_id
)
}
}
impl error::Error for SetActualSpanIdError {}