Skip to main content

ruststream_lapin/testing/
subscriber.rs

1//! The in-process subscriber and its delivery type.
2
3use std::sync::Arc;
4
5use futures::Stream;
6use ruststream::testing::Coordinator;
7use ruststream::{AckError, Headers, IncomingMessage, Partitioned, Subscriber};
8
9use super::broker::TestBrokerState;
10use super::router::{DeliveryReceiver, DeliverySender, SubscriptionId, TestDelivery};
11use crate::error::AmqpError;
12
13/// In-process subscriber on one queue name.
14///
15/// Yielded messages settle like the real transport: ack finalizes, `nack(true)` re-enqueues to
16/// this same subscription, `nack(false)` drops.
17pub struct LapinTestSubscriber {
18    state: Arc<TestBrokerState>,
19    id: SubscriptionId,
20    queue: String,
21    sender: DeliverySender,
22    receiver: DeliveryReceiver,
23    coordinator: Option<Coordinator>,
24}
25
26impl LapinTestSubscriber {
27    pub(crate) fn open(state: &Arc<TestBrokerState>, queue: String) -> Self {
28        let (id, sender, receiver) = state.router.subscribe(queue.clone());
29        let coordinator = state.coordinator();
30        Self {
31            state: Arc::clone(state),
32            id,
33            queue,
34            sender,
35            receiver,
36            coordinator,
37        }
38    }
39
40    /// The queue this subscriber consumes from.
41    #[must_use]
42    pub fn queue(&self) -> &str {
43        &self.queue
44    }
45}
46
47impl Drop for LapinTestSubscriber {
48    fn drop(&mut self) {
49        self.state.router.unsubscribe(self.id);
50    }
51}
52
53impl std::fmt::Debug for LapinTestSubscriber {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        f.debug_struct("LapinTestSubscriber")
56            .field("queue", &self.queue)
57            .finish_non_exhaustive()
58    }
59}
60
61impl Subscriber for LapinTestSubscriber {
62    type Message = LapinTestMessage;
63    type Error = AmqpError;
64
65    /// Streams injected deliveries; never yields an error.
66    ///
67    /// # Cancel safety
68    ///
69    /// Cancel safe and re-enterable: the receiver is polled in place, so dropping the returned
70    /// stream loses nothing and `stream` can be called again.
71    fn stream(&mut self) -> impl Stream<Item = Result<Self::Message, Self::Error>> + Send + '_ {
72        let Self {
73            receiver,
74            sender,
75            coordinator,
76            ..
77        } = self;
78        futures::stream::poll_fn(move |cx| {
79            receiver.poll_recv(cx).map(|delivery| {
80                delivery.map(|delivery| {
81                    Ok(LapinTestMessage {
82                        delivery: Some(delivery),
83                        sender: sender.clone(),
84                        coordinator: coordinator.clone(),
85                    })
86                })
87            })
88        })
89    }
90}
91
92/// One in-process delivery.
93pub struct LapinTestMessage {
94    delivery: Option<TestDelivery>,
95    sender: DeliverySender,
96    coordinator: Option<Coordinator>,
97}
98
99impl LapinTestMessage {
100    fn take(&mut self) -> TestDelivery {
101        // The settle methods consume `self`, so a second settle cannot compile; reaching this
102        // twice is an internal invariant violation.
103        self.delivery
104            .take()
105            .expect("LapinTestMessage settled twice")
106    }
107}
108
109impl Drop for LapinTestMessage {
110    fn drop(&mut self) {
111        // Balance the router's `enqueued` exactly once per delivery, whatever the dispatch
112        // path did (ack, nack, panic, or plain drop).
113        if let Some(coordinator) = self.coordinator.take() {
114            coordinator.consumed();
115        }
116    }
117}
118
119impl std::fmt::Debug for LapinTestMessage {
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        f.debug_struct("LapinTestMessage")
122            .field("delivery", &self.delivery)
123            .finish_non_exhaustive()
124    }
125}
126
127impl IncomingMessage for LapinTestMessage {
128    fn payload(&self) -> &[u8] {
129        &self
130            .delivery
131            .as_ref()
132            .expect("message accessed after settlement")
133            .payload
134    }
135
136    fn headers(&self) -> &Headers {
137        &self
138            .delivery
139            .as_ref()
140            .expect("message accessed after settlement")
141            .headers
142    }
143
144    /// The partition key from the `PARTITION_KEY_HEADER`, mirroring the real message so keyed
145    /// worker lanes behave the same in-process.
146    fn partition_key(&self) -> Option<&[u8]> {
147        self.headers().get(crate::PARTITION_KEY_HEADER)
148    }
149
150    /// Finalizes the delivery.
151    ///
152    /// # Errors
153    ///
154    /// Never fails; the in-process transport has no channel to lose.
155    async fn ack(mut self) -> Result<(), AckError> {
156        drop(self.take());
157        Ok(())
158    }
159
160    /// Re-enqueues to the same subscription (`requeue = true`) or drops (`requeue = false`).
161    ///
162    /// # Errors
163    ///
164    /// Never fails; the in-process transport has no channel to lose.
165    async fn nack(mut self, requeue: bool) -> Result<(), AckError> {
166        let delivery = self.take();
167        if requeue && self.sender.send(delivery).is_ok() {
168            // This bypasses the router fanout, so account for the new in-flight delivery here.
169            if let Some(coordinator) = &self.coordinator {
170                coordinator.enqueued();
171            }
172        }
173        Ok(())
174    }
175}
176
177impl Partitioned for LapinTestMessage {
178    /// The partition key from the `PARTITION_KEY_HEADER`, mirroring the real message.
179    fn partition_key(&self) -> Option<&[u8]> {
180        self.headers().get(crate::PARTITION_KEY_HEADER)
181    }
182}