use std::any::{TypeId, type_name};
use std::fmt;
use std::sync::Arc;
use std::time::{Duration, Instant as WallClock};
use thiserror::Error;
use tokio::task::JoinHandle;
use tokio_util::task::TaskTracker;
use crate::OutgoingDestination;
use crate::runtime::{
App, ConnectedLifecycle, HeadersUnset, LifecycleHook, OutSlot, PublishBuilder,
RegisteredBroker, RustStreamError, Shutdown, Starter, TestParts,
};
use crate::runtime::{MessageBody, UnnamedCodec, message_of};
use crate::{Broker, DefaultPublish, PublishPolicy, Publisher};
use super::assertions::PublishedAssertions;
use super::broker::{LivePublish, TestableBroker, TestableRegistration, live_publish};
use super::coordinator::{Coordinator, LiveSubscription};
use super::handle::{BrokerHandle, Harness, InjectSink, Target, Transport};
const DEFAULT_MAX_STEPS: usize = 10_000;
const DEFAULT_SETTLE_DEADLINE: Duration = Duration::from_secs(10);
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum TestError {
#[error("startup hook failed: {0}")]
Startup(#[source] Box<dyn std::error::Error + Send + Sync>),
#[error("subscription failed: {0}")]
Subscribe(#[source] Box<dyn std::error::Error + Send + Sync>),
#[error("the reaction did not settle within {processed} dispatched deliveries")]
NotQuiescent {
processed: usize,
},
#[error(
"the live reaction did not settle within {deadline:?}: {}",
unsettled(subscription.as_deref(), *handled, *expected)
)]
NotSettled {
subscription: Option<String>,
handled: usize,
expected: usize,
deadline: Duration,
},
#[error("broker {broker} failed to connect: {source}")]
Connect {
broker: String,
#[source]
source: Box<dyn std::error::Error + Send + Sync>,
},
#[error("broker {broker} refused the test's publish: {source}")]
Publish {
broker: String,
#[source]
source: Box<dyn std::error::Error + Send + Sync>,
},
#[error("publish after the service shut down")]
ShutDown,
#[error("more than one broker is registered; address one with broker::<B>() or broker_named()")]
Ambiguous,
#[error(
"broker {0} has no in-process mode; enable the `testing` feature of the crate that \
provides it in `[dev-dependencies]`"
)]
NoTransport(String),
#[error(
"the clock is paused, and a live test runs against brokers on the real clock; drop \
`start_paused` / `tokio::time::pause()` from this test, or start it in process with \
`TestApp::start`"
)]
PausedClock,
#[error("failed to encode the message: {0}")]
Encode(String),
}
fn unsettled(subscription: Option<&str>, handled: usize, expected: usize) -> String {
subscription.map_or_else(
|| "a handler was still running".to_owned(),
|name| format!("subscription {name:?} handled {handled} of {expected} deliveries"),
)
}
pub(super) struct BrokerEntry {
label: Option<String>,
lifecycle: Box<dyn ConnectedLifecycle>,
registration: Option<&'static TestableRegistration>,
}
impl BrokerEntry {
fn testable(&self) -> Option<&dyn TestableBroker> {
self.registration
.and_then(|registration| registration.resolve(self.lifecycle.as_any()))
}
fn display(&self) -> String {
self.label
.clone()
.unwrap_or_else(|| self.lifecycle.name().to_owned())
}
}
fn same_names(destination: &str, subscriptions: &[&str]) -> Vec<usize> {
subscriptions
.iter()
.enumerate()
.filter(|(_, name)| **name == destination)
.map(|(position, _)| position)
.collect()
}
fn registration_of(broker: TypeId) -> Option<&'static TestableRegistration> {
inventory::iter::<TestableRegistration>
.into_iter()
.find(|registration| registration.registers(broker))
}
pub struct TestBrokers<'a> {
entries: &'a [BrokerEntry],
}
impl fmt::Debug for TestBrokers<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TestBrokers")
.field("brokers", &self.entries.len())
.finish_non_exhaustive()
}
}
impl TestBrokers<'_> {
#[must_use]
pub fn broker<B: Broker + 'static>(&self) -> &B::Connected {
let mut found = self
.entries
.iter()
.filter(|e| e.lifecycle.broker_type() == TypeId::of::<B>())
.filter_map(|e| e.lifecycle.as_any().downcast_ref::<B::Connected>());
let first = found
.next()
.unwrap_or_else(|| panic!("no registered broker of type {}", type_name::<B>()));
assert!(
found.next().is_none(),
"more than one broker of type {} is registered",
type_name::<B>(),
);
first
}
}
pub struct TestApp<State> {
entries: Vec<BrokerEntry>,
mode: Mode,
coordinator: Coordinator,
#[allow(dead_code)]
state: Arc<State>,
shutdown: Shutdown,
handles: Vec<JoinHandle<()>>,
continuations: TaskTracker,
shutdown_timeout: Option<Duration>,
}
impl<State> fmt::Debug for TestApp<State> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TestApp")
.field("mode", &self.mode)
.field("brokers", &self.entries.len())
.field("subscribers", &self.handles.len())
.finish_non_exhaustive()
}
}
#[derive(Debug)]
pub(super) enum Mode {
InProcess,
Live {
deadline: Duration,
subscriptions: Vec<LiveSubscription>,
},
}
impl Mode {
pub(super) async fn settle(
&self,
coordinator: &Coordinator,
brokers: &[BrokerEntry],
) -> Result<(), TestError> {
match self {
Self::InProcess => coordinator.drive().await,
Self::Live {
deadline,
subscriptions,
} => {
let routing = |scope_id: usize, destination: &str, names: &[&str]| {
brokers[scope_id].testable().map_or_else(
|| same_names(destination, names),
|broker| broker.routes(destination, names),
)
};
coordinator
.settle_live(subscriptions, &routing, *deadline)
.await
}
}
}
}
fn clock_is_paused() -> bool {
let before = tokio::time::Instant::now();
let wall = WallClock::now();
while wall.elapsed() < Duration::from_micros(50) {
std::hint::spin_loop();
}
tokio::time::Instant::now() == before
}
impl<State: Send + Sync + 'static> TestApp<State> {
pub async fn start<A>(app: A) -> Result<Self, TestError>
where
A: App<State = State>,
{
let (coordinator, parts) = Self::prepare(app);
let TestParts {
brokers,
starters,
state_init,
after_startup,
shutdown_timeout,
continuations,
..
} = parts;
let registered = registrations(brokers)?;
let state = state_init().await.map_err(TestError::Startup)?;
let entries = connect_in_process(registered, &coordinator).await?;
Self::spawn(SpawnArgs {
coordinator,
entries,
mode: Mode::InProcess,
starters,
after_startup,
continuations,
shutdown_timeout,
state: Arc::new(state),
})
.await
}
pub async fn start_live<A>(app: A) -> Result<Self, TestError>
where
A: App<State = State>,
{
Self::start_live_within(app, DEFAULT_SETTLE_DEADLINE).await
}
pub async fn start_live_within<A>(app: A, deadline: Duration) -> Result<Self, TestError>
where
A: App<State = State>,
{
if clock_is_paused() {
return Err(TestError::PausedClock);
}
let (coordinator, parts) = Self::prepare(app);
let TestParts {
brokers,
starters,
state_init,
after_startup,
shutdown_timeout,
continuations,
test_hooks,
} = parts;
let state = state_init().await.map_err(TestError::Startup)?;
let entries = connect_live(brokers).await?;
let subscriptions = test_hooks.subscriptions();
let mode = Mode::Live {
deadline,
subscriptions,
};
Self::spawn(SpawnArgs {
coordinator,
entries,
mode,
starters,
after_startup,
continuations,
shutdown_timeout,
state: Arc::new(state),
})
.await
}
pub async fn with_state<A, F>(app: A, build: F) -> Result<Self, TestError>
where
A: App<State = State>,
F: FnOnce(&TestBrokers<'_>) -> State,
{
let (coordinator, parts) = Self::prepare(app);
let TestParts {
brokers,
starters,
after_startup,
shutdown_timeout,
continuations,
..
} = parts;
let registered = registrations(brokers)?;
let entries = connect_in_process(registered, &coordinator).await?;
let state = build(&TestBrokers { entries: &entries });
Self::spawn(SpawnArgs {
coordinator,
entries,
mode: Mode::InProcess,
starters,
after_startup,
continuations,
shutdown_timeout,
state: Arc::new(state),
})
.await
}
fn prepare<A>(app: A) -> (Coordinator, TestParts<State>)
where
A: App<State = State>,
{
let parts = app.into_test_parts();
let coordinator = Coordinator::new(DEFAULT_MAX_STEPS);
parts.test_hooks.install(coordinator.clone());
(coordinator, parts)
}
async fn spawn(args: SpawnArgs<State>) -> Result<Self, TestError> {
let SpawnArgs {
coordinator,
entries,
mode,
starters,
after_startup,
continuations,
shutdown_timeout,
state,
} = args;
coordinator.locate(entries.iter().map(|entry| entry.lifecycle.as_any()));
let shutdown = Shutdown::new();
let mut handles = Vec::with_capacity(starters.len());
let failed = 'start: {
for starter in starters {
match starter(state.clone(), shutdown.clone()).await {
Ok(handle) => handles.push(handle),
Err(err) => break 'start Some(TestError::Subscribe(err)),
}
}
for hook in after_startup {
if let Err(err) = hook(state.clone()).await {
break 'start Some(TestError::Startup(err));
}
}
None
};
if let Some(err) = failed {
stop_dispatch(&shutdown, handles, shutdown_timeout).await;
drain_continuations(&continuations, shutdown_timeout).await;
shut_down(entries).await;
return Err(err);
}
Ok(Self {
entries,
mode,
coordinator,
state,
shutdown,
handles,
continuations,
shutdown_timeout,
})
}
#[must_use]
pub fn broker<B>(&self) -> BrokerHandle<'_>
where
B: Broker + 'static,
B::Connected: DefaultPublish,
<<B::Connected as DefaultPublish>::Policy as PublishPolicy<B::Connected>>::Live: Publisher,
{
let mut matches = self
.entries
.iter()
.filter(|e| e.lifecycle.broker_type() == TypeId::of::<B>());
let first = matches
.next()
.unwrap_or_else(|| panic!("no registered broker of type {}", type_name::<B>()));
assert!(
matches.next().is_none(),
"more than one broker of type {} is registered; address one with broker_named(label)",
type_name::<B>(),
);
self.handle(first, Some(live_publish::<B>))
}
#[must_use]
pub fn broker_named(&self, label: &str) -> BrokerHandle<'_> {
let entry = self
.entries
.iter()
.find(|e| e.label.as_deref() == Some(label))
.unwrap_or_else(|| panic!("no broker labeled {label:?}"));
self.handle(entry, None)
}
#[must_use]
pub fn out<M: OutSlot>(&self) -> PublishedAssertions<()> {
PublishedAssertions::captured(
format!("Out slot `{}`", M::NAME),
self.coordinator.slot_published(M::NAME),
)
}
fn handle<'a>(
&'a self,
entry: &'a BrokerEntry,
typed: Option<LivePublish>,
) -> BrokerHandle<'a> {
let scope_id = self
.entries
.iter()
.position(|e| std::ptr::eq(e, entry))
.expect("entry belongs to this app");
let transport = match &self.mode {
Mode::InProcess => Transport::InProcess(
entry
.testable()
.expect("an in-process broker connected through its registration"),
),
Mode::Live { .. } => Transport::Live {
connected: entry.lifecycle.as_any(),
publish: typed.or_else(|| entry.registration.map(TestableRegistration::live)),
},
};
BrokerHandle {
harness: Harness {
coordinator: &self.coordinator,
shutdown: &self.shutdown,
mode: &self.mode,
brokers: &self.entries,
},
scope_id,
transport,
label: entry.display(),
}
}
#[cfg(any(feature = "json", feature = "cbor", feature = "msgpack"))]
pub async fn publish<T: serde::Serialize + Sync>(
&self,
name: &str,
value: &T,
) -> Result<(), TestError> {
if self.entries.len() != 1 {
return Err(TestError::Ambiguous);
}
self.handle(&self.entries[0], None)
.publish(name, value)
.await
}
pub fn message<'a, T>(
&'a self,
value: &'a T,
) -> PublishBuilder<InjectSink<'a>, MessageBody<'a, T>, UnnamedCodec, HeadersUnset, T::Form>
where
T: OutgoingDestination,
{
message_of(self.sole_sink(), value, UnnamedCodec::new())
}
fn sole_sink(&self) -> InjectSink<'_> {
match self.entries.as_slice() {
[only] => self.handle(only, None).sink(),
_ => InjectSink(Target::Ambiguous),
}
}
pub async fn settle(&self) -> Result<(), TestError> {
self.mode.settle(&self.coordinator, &self.entries).await
}
pub async fn advance(&self, by: Duration) -> Result<(), TestError> {
match self.mode {
Mode::InProcess => tokio::time::advance(by).await,
Mode::Live { .. } => tokio::time::sleep(by).await,
}
self.coordinator.fire_due_timers().await;
self.settle().await
}
pub async fn drain(&self) {
while !self.continuations.is_empty() {
tokio::task::yield_now().await;
}
}
pub fn run_result(&self) -> Result<(), RustStreamError> {
self.shutdown
.peek_failure()
.map_or(Ok(()), |reason| Err(RustStreamError::Dispatch(reason)))
}
pub fn assert_running(&self) {
assert!(
!self.shutdown.is_cancelled(),
"expected the service to be running, but it was shut down: {:?}",
self.shutdown.peek_failure(),
);
}
pub fn assert_shut_down(&self) {
assert!(
self.shutdown.is_cancelled(),
"expected the service to be shut down, but it was still running",
);
}
pub async fn shutdown(self) -> Result<(), RustStreamError> {
stop_dispatch(&self.shutdown, self.handles, self.shutdown_timeout).await;
self.continuations.close();
self.continuations.wait().await;
shut_down(self.entries).await;
self.shutdown
.taken_failure()
.map_or(Ok(()), |reason| Err(RustStreamError::Dispatch(reason)))
}
}
async fn stop_dispatch(
shutdown: &Shutdown,
handles: Vec<JoinHandle<()>>,
timeout: Option<Duration>,
) {
shutdown.cancel();
let until = timeout.map(|timeout| tokio::time::Instant::now() + timeout);
for mut handle in handles {
let Some(until) = until else {
let _ = handle.await;
continue;
};
if tokio::time::timeout_at(until, &mut handle).await.is_err() {
handle.abort();
let _ = handle.await;
}
}
}
async fn drain_continuations(continuations: &TaskTracker, timeout: Option<Duration>) {
continuations.close();
match timeout {
Some(timeout) => {
let _ = tokio::time::timeout(timeout, continuations.wait()).await;
}
None => continuations.wait().await,
}
}
async fn shut_down(entries: Vec<BrokerEntry>) {
for entry in entries.into_iter().rev() {
let broker = entry.display();
if let Err(err) = entry.lifecycle.shutdown().await {
tracing::debug!(
target: "ruststream::testing",
broker = %broker,
error = %err,
"broker shutdown after the test failed",
);
}
}
}
fn registrations(
brokers: Vec<RegisteredBroker>,
) -> Result<Vec<(RegisteredBroker, &'static TestableRegistration)>, TestError> {
brokers
.into_iter()
.map(|broker| {
let registration = registration_of(broker.lifecycle.broker_type())
.ok_or_else(|| TestError::NoTransport(broker.lifecycle.broker_name().to_owned()))?;
Ok((broker, registration))
})
.collect()
}
async fn connect_in_process(
registered: Vec<(RegisteredBroker, &'static TestableRegistration)>,
coordinator: &Coordinator,
) -> Result<Vec<BrokerEntry>, TestError> {
let mut entries = Vec::with_capacity(registered.len());
for (RegisteredBroker { lifecycle, label }, registration) in registered {
let broker = label
.clone()
.unwrap_or_else(|| lifecycle.broker_name().to_owned());
let lifecycle = match lifecycle
.connect_in_process(registration.transition())
.await
{
Ok(lifecycle) => lifecycle,
Err(source) => {
shut_down(entries).await;
return Err(TestError::Connect { broker, source });
}
};
let entry = BrokerEntry {
label,
lifecycle,
registration: Some(registration),
};
entry
.testable()
.expect("a registration resolves the connected form its own transition produced")
.install_coordinator(coordinator.clone());
entries.push(entry);
}
Ok(entries)
}
async fn connect_live(brokers: Vec<RegisteredBroker>) -> Result<Vec<BrokerEntry>, TestError> {
let mut entries = Vec::with_capacity(brokers.len());
for RegisteredBroker { lifecycle, label } in brokers {
let broker = label
.clone()
.unwrap_or_else(|| lifecycle.broker_name().to_owned());
let registration = registration_of(lifecycle.broker_type());
let lifecycle = match lifecycle.connect().await {
Ok(lifecycle) => lifecycle,
Err(source) => {
shut_down(entries).await;
return Err(TestError::Connect { broker, source });
}
};
entries.push(BrokerEntry {
label,
lifecycle,
registration,
});
}
Ok(entries)
}
struct SpawnArgs<State> {
coordinator: Coordinator,
entries: Vec<BrokerEntry>,
mode: Mode,
starters: Vec<Starter<State>>,
after_startup: Vec<LifecycleHook<State>>,
continuations: TaskTracker,
shutdown_timeout: Option<Duration>,
state: Arc<State>,
}