use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, OnceLock};
use bytes::Bytes;
use ruststream::testing::{Coordinator, TestableBroker};
use ruststream::{
Broker, ConnectedBroker, DefaultPublish, DescribeServer, OutgoingMessage, RawMessage,
ServerSpec, Subscribe,
};
use super::publisher::{LapinTestPublish, LapinTestPublisher};
use super::router::KeyRouter;
use super::subscriber::LapinTestSubscriber;
use crate::error::AmqpError;
pub(crate) struct TestBrokerState {
pub(crate) router: KeyRouter,
closed: AtomicBool,
coordinator: OnceLock<Coordinator>,
}
impl TestBrokerState {
pub(crate) fn install(&self, coordinator: Coordinator) {
let _ = self.coordinator.set(coordinator);
}
pub(crate) fn coordinator(&self) -> Option<Coordinator> {
self.coordinator.get().cloned()
}
pub(crate) fn ensure_live(&self, target: &str) -> Result<(), AmqpError> {
if self.closed.load(Ordering::Acquire) {
return Err(AmqpError::closed(target));
}
Ok(())
}
}
impl Default for TestBrokerState {
fn default() -> Self {
Self {
router: KeyRouter::default(),
closed: AtomicBool::new(false),
coordinator: OnceLock::new(),
}
}
}
impl std::fmt::Debug for TestBrokerState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TestBrokerState")
.field("router", &self.router)
.field("closed", &self.closed.load(Ordering::Relaxed))
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone, Default)]
#[must_use]
pub struct LapinTestBroker {
state: Arc<TestBrokerState>,
}
impl LapinTestBroker {
pub fn new() -> Self {
Self::default()
}
}
impl Broker for LapinTestBroker {
type Error = AmqpError;
type Connected = ConnectedLapinTestBroker;
async fn connect(self) -> Result<Self::Connected, Self::Error> {
Ok(ConnectedLapinTestBroker { state: self.state })
}
}
impl DescribeServer for LapinTestBroker {
fn describe_server(&self) -> ServerSpec {
ServerSpec::in_process("amqp")
}
}
#[derive(Debug, Clone)]
pub struct ConnectedLapinTestBroker {
state: Arc<TestBrokerState>,
}
impl ConnectedLapinTestBroker {
pub(crate) fn state(&self) -> Arc<TestBrokerState> {
Arc::clone(&self.state)
}
#[allow(clippy::unused_async)]
pub async fn subscribe(
&self,
queue: impl Into<String>,
) -> Result<LapinTestSubscriber, AmqpError> {
let queue = queue.into();
if queue.is_empty() {
return Err(AmqpError::InvalidOptions(
"queue name must not be empty; subscribe with the queue the handler consumes"
.to_owned(),
));
}
self.state.ensure_live(&queue)?;
Ok(LapinTestSubscriber::open(&self.state, queue))
}
#[must_use]
pub fn publisher(&self, policy: LapinTestPublish) -> LapinTestPublisher {
policy.bind(self)
}
}
impl ConnectedBroker for ConnectedLapinTestBroker {
type Error = AmqpError;
type Closed = ();
async fn shutdown(self) -> Result<Self::Closed, Self::Error> {
self.state.closed.store(true, Ordering::Release);
self.state.router.clear();
Ok(())
}
}
#[allow(clippy::use_self)]
impl Subscribe for ConnectedLapinTestBroker {
type Subscriber = LapinTestSubscriber;
async fn subscribe(&self, name: &str) -> Result<Self::Subscriber, Self::Error> {
ConnectedLapinTestBroker::subscribe(self, name).await
}
}
impl DefaultPublish for ConnectedLapinTestBroker {
type Policy = LapinTestPublish;
}
impl TestableBroker for ConnectedLapinTestBroker {
fn install_coordinator(&self, coordinator: Coordinator) {
self.state.install(coordinator);
}
fn inject(&self, message: OutgoingMessage<'_>) {
self.state.router.publish(
message.name(),
&Bytes::copy_from_slice(message.payload()),
message.headers(),
self.state.coordinator().as_ref(),
);
}
fn published(&self, name: &str) -> Vec<RawMessage> {
self.state.router.published(name)
}
}
ruststream::register_testable_broker!(ConnectedLapinTestBroker);