ruststream_lapin/testing/
subscriber.rs1use 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
13pub 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 #[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 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
92pub struct LapinTestMessage {
94 delivery: Option<TestDelivery>,
95 sender: DeliverySender,
96 coordinator: Option<Coordinator>,
97}
98
99impl LapinTestMessage {
100 fn take(&mut self) -> TestDelivery {
101 self.delivery
104 .take()
105 .expect("LapinTestMessage settled twice")
106 }
107}
108
109impl Drop for LapinTestMessage {
110 fn drop(&mut self) {
111 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 fn partition_key(&self) -> Option<&[u8]> {
147 self.headers().get(crate::PARTITION_KEY_HEADER)
148 }
149
150 async fn ack(mut self) -> Result<(), AckError> {
156 drop(self.take());
157 Ok(())
158 }
159
160 async fn nack(mut self, requeue: bool) -> Result<(), AckError> {
166 let delivery = self.take();
167 if requeue && self.sender.send(delivery).is_ok() {
168 if let Some(coordinator) = &self.coordinator {
170 coordinator.enqueued();
171 }
172 }
173 Ok(())
174 }
175}
176
177impl Partitioned for LapinTestMessage {
178 fn partition_key(&self) -> Option<&[u8]> {
180 self.headers().get(crate::PARTITION_KEY_HEADER)
181 }
182}