allama_http/
replay.rs

1use std::collections::VecDeque;
2use std::fmt;
3use std::sync::{Arc, Mutex};
4
5use bytes::Bytes;
6use futures_util::StreamExt;
7
8use crate::{
9    Body, Headers, HttpRequest, HttpResponse, HttpTransport, SendFuture, TransportError,
10    TransportErrorKind,
11};
12
13/// An [`HttpTransport`] that serves canned responses in order and records
14/// every request.
15///
16/// Use it to run a real provider adapter (request encoding, error
17/// classification, SSE decoding) against recorded wire data without a
18/// server: in unit tests, on `wasm32`, or to replay a captured exchange.
19/// When the queue runs out, requests fail with a
20/// [`Connect`](TransportErrorKind::Connect) error.
21///
22/// Clone is cheap and clones share the queue and the recorded requests.
23#[derive(Clone, Default)]
24pub struct ReplayTransport {
25    state: Arc<Mutex<State>>,
26}
27
28#[derive(Default)]
29struct State {
30    responses: VecDeque<ReplayResponse>,
31    requests: Vec<HttpRequest>,
32}
33
34impl fmt::Debug for ReplayTransport {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        let state = self.lock();
37        f.debug_struct("ReplayTransport")
38            .field("remaining", &state.responses.len())
39            .field("requests", &state.requests.len())
40            .finish()
41    }
42}
43
44impl ReplayTransport {
45    /// A transport serving `responses` in order.
46    #[must_use]
47    pub fn new(responses: impl IntoIterator<Item = ReplayResponse>) -> Self {
48        Self {
49            state: Arc::new(Mutex::new(State {
50                responses: responses.into_iter().collect(),
51                requests: Vec::new(),
52            })),
53        }
54    }
55
56    /// Queue another response.
57    pub fn push(&self, response: ReplayResponse) {
58        self.lock().responses.push_back(response);
59    }
60
61    /// The requests sent so far, in order.
62    #[must_use]
63    pub fn requests(&self) -> Vec<HttpRequest> {
64        self.lock().requests.clone()
65    }
66
67    /// How many queued responses are left.
68    #[must_use]
69    pub fn remaining(&self) -> usize {
70        self.lock().responses.len()
71    }
72
73    fn lock(&self) -> std::sync::MutexGuard<'_, State> {
74        self.state.lock().unwrap_or_else(|e| e.into_inner())
75    }
76}
77
78impl HttpTransport for ReplayTransport {
79    fn send(&self, request: HttpRequest) -> SendFuture<'_> {
80        let next = {
81            let mut state = self.lock();
82            state.requests.push(request);
83            state.responses.pop_front()
84        };
85        Box::pin(async move {
86            let response = next.ok_or_else(|| {
87                TransportError::new(TransportErrorKind::Connect, "no replayed response left")
88            })?;
89            if let Some((kind, message)) = response.failure {
90                return Err(TransportError::new(kind, message));
91            }
92            let chunks = futures_util::stream::iter(response.chunks.into_iter().map(Ok));
93            let body: Body = if response.stalls {
94                Box::pin(chunks.chain(futures_util::stream::pending()))
95            } else {
96                Box::pin(chunks)
97            };
98            Ok(HttpResponse::new(response.status, response.headers, body))
99        })
100    }
101}
102
103/// A canned response served by [`ReplayTransport`].
104///
105/// The body arrives in the chunks it was built with: one for
106/// [`new`](Self::new) and [`json`](Self::json), one per event for
107/// [`sse`](Self::sse), the given ones for [`chunks`](Self::chunks).
108/// [`split_every`](Self::split_every) re-chunks any body, down to a byte at
109/// a time, to test that a decoder does not depend on where the network
110/// splits it.
111#[derive(Debug, Clone)]
112pub struct ReplayResponse {
113    status: u16,
114    headers: Headers,
115    chunks: Vec<Bytes>,
116    failure: Option<(TransportErrorKind, String)>,
117    stalls: bool,
118}
119
120impl ReplayResponse {
121    /// A response with `body` in one chunk.
122    #[must_use]
123    pub fn new(status: u16, body: impl Into<Bytes>) -> Self {
124        Self::chunks(status, [body])
125    }
126
127    /// A response whose body arrives in `chunks`, in order: a stream body
128    /// in another framing than SSE (NDJSON lines), or chunk boundaries a
129    /// test places on purpose.
130    ///
131    /// ```
132    /// use allama_http::{HttpRequest, HttpTransport, ReplayResponse, ReplayTransport};
133    /// use futures_util::StreamExt;
134    ///
135    /// # async fn example() {
136    /// let ndjson = ReplayResponse::chunks(200, ["{\"text\":\"Hel\"}\n{\"te", "xt\":\"lo\"}\n"])
137    ///     .header("content-type", "application/x-ndjson");
138    /// let transport = ReplayTransport::new([ndjson]);
139    /// let response = transport.send(HttpRequest::get("https://x")).await.unwrap();
140    /// assert_eq!(response.into_body().count().await, 2);
141    /// # }
142    /// ```
143    #[must_use]
144    pub fn chunks(status: u16, chunks: impl IntoIterator<Item = impl Into<Bytes>>) -> Self {
145        Self {
146            status,
147            headers: Headers::new(),
148            chunks: chunks.into_iter().map(Into::into).collect(),
149            failure: None,
150            stalls: false,
151        }
152    }
153
154    /// A JSON response with `body` as written: a captured body, or one
155    /// that is not valid JSON on purpose.
156    #[must_use]
157    pub fn json(status: u16, body: impl Into<Bytes>) -> Self {
158        Self::new(status, body).header("content-type", "application/json")
159    }
160
161    /// A JSON response with `value` serialized as its body.
162    ///
163    /// # Panics
164    ///
165    /// When `value` does not serialize.
166    #[must_use]
167    pub fn json_value(status: u16, value: &impl serde::Serialize) -> Self {
168        let body = serde_json::to_vec(value).expect("serialize replayed JSON response");
169        Self::json(status, body)
170    }
171
172    /// A `200` server-sent-events response, delivered one event per chunk
173    /// the way a network would split it.
174    #[must_use]
175    pub fn sse(body: impl AsRef<str>) -> Self {
176        let chunks = body
177            .as_ref()
178            .split_inclusive("\n\n")
179            .map(|event| Bytes::from(event.to_owned()));
180        Self::chunks(200, chunks).header("content-type", "text/event-stream")
181    }
182
183    /// A `200` newline-delimited JSON response, delivered one line per
184    /// chunk the way a network would split it: the NDJSON counterpart of
185    /// [`sse`](Self::sse). Use [`chunks`](Self::chunks) to place chunk
186    /// boundaries yourself, or another status.
187    ///
188    /// ```
189    /// use allama_http::{HttpRequest, HttpTransport, ReplayResponse, ReplayTransport};
190    /// use futures_util::StreamExt;
191    ///
192    /// # async fn example() {
193    /// let transport = ReplayTransport::new([ReplayResponse::ndjson(
194    ///     "{\"delta\":\"Hel\"}\n{\"delta\":\"lo\"}\n{\"done\":true}\n",
195    /// )]);
196    /// let response = transport.send(HttpRequest::get("https://x")).await.unwrap();
197    /// assert_eq!(response.headers.get("content-type"), Some("application/x-ndjson"));
198    /// assert_eq!(response.into_body().count().await, 3);
199    /// # }
200    /// ```
201    #[must_use]
202    pub fn ndjson(body: impl AsRef<str>) -> Self {
203        let chunks = body
204            .as_ref()
205            .split_inclusive('\n')
206            .map(|line| Bytes::from(line.to_owned()));
207        Self::chunks(200, chunks).header("content-type", "application/x-ndjson")
208    }
209
210    /// A request that gets no response, failing with `kind`.
211    #[must_use]
212    pub fn failure(kind: TransportErrorKind, message: impl Into<String>) -> Self {
213        Self {
214            status: 0,
215            headers: Headers::new(),
216            chunks: Vec::new(),
217            failure: Some((kind, message.into())),
218            stalls: false,
219        }
220    }
221
222    /// Add a response header.
223    #[must_use]
224    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
225        self.headers.insert(name, value);
226        self
227    }
228
229    /// Deliver the body in chunks of `size` bytes (the last one may be
230    /// shorter), whatever chunks it was built with: `split_every(1)` sends
231    /// one byte at a time, splitting every line ending and UTF-8 sequence.
232    ///
233    /// ```
234    /// use allama_http::{HttpRequest, HttpTransport, ReplayResponse, ReplayTransport};
235    /// use futures_util::StreamExt;
236    ///
237    /// # async fn example() {
238    /// let transport = ReplayTransport::new([ReplayResponse::sse("data: hi\n\n").split_every(4)]);
239    /// let response = transport.send(HttpRequest::get("https://x")).await.unwrap();
240    /// assert_eq!(response.into_body().count().await, 3);
241    /// # }
242    /// ```
243    ///
244    /// # Panics
245    ///
246    /// When `size` is 0.
247    #[must_use]
248    pub fn split_every(mut self, size: usize) -> Self {
249        assert!(
250            size > 0,
251            "a replayed body cannot be split into empty chunks"
252        );
253        let body: Vec<u8> = self.chunks.concat();
254        self.chunks = body.chunks(size).map(Bytes::copy_from_slice).collect();
255        self
256    }
257
258    /// Never end the body after its chunks, like a connection that stalls
259    /// mid-stream: for testing idle timeouts such as [`StreamTimeouts`].
260    ///
261    /// [`StreamTimeouts`]: crate::StreamTimeouts
262    #[must_use]
263    pub fn then_stall(mut self) -> Self {
264        self.stalls = true;
265        self
266    }
267}
268
269#[cfg(all(test, not(target_arch = "wasm32")))]
270mod tests {
271    use futures_util::StreamExt;
272
273    use super::*;
274
275    #[tokio::test]
276    async fn responses_are_served_in_order_and_requests_recorded() {
277        let transport = ReplayTransport::new([
278            ReplayResponse::json(200, r#"{"n":1}"#).header("request-id", "req_1"),
279            ReplayResponse::failure(TransportErrorKind::Timeout, "slow"),
280        ]);
281        let first = transport
282            .send(HttpRequest::post("https://x/a", "1"))
283            .await
284            .unwrap();
285        assert_eq!(first.headers.get("Request-Id"), Some("req_1"));
286        assert_eq!(first.text().await.unwrap(), r#"{"n":1}"#);
287        let second = transport
288            .send(HttpRequest::get("https://x/b"))
289            .await
290            .unwrap_err();
291        assert_eq!(second.kind(), TransportErrorKind::Timeout);
292        let exhausted = transport
293            .send(HttpRequest::get("https://x/c"))
294            .await
295            .unwrap_err();
296        assert_eq!(exhausted.kind(), TransportErrorKind::Connect);
297
298        let urls: Vec<_> = transport.requests().into_iter().map(|r| r.url).collect();
299        assert_eq!(urls, ["https://x/a", "https://x/b", "https://x/c"]);
300        assert_eq!(transport.remaining(), 0);
301    }
302
303    #[tokio::test]
304    async fn sse_is_chunked_per_event() {
305        let transport = ReplayTransport::default();
306        transport.push(ReplayResponse::sse("data: a\n\ndata: b\n\n"));
307        let response = transport.send(HttpRequest::get("https://x")).await.unwrap();
308        assert_eq!(
309            response.headers.get("content-type"),
310            Some("text/event-stream")
311        );
312        let chunks: Vec<_> = response.into_body().collect().await;
313        assert_eq!(chunks.len(), 2);
314    }
315
316    #[tokio::test]
317    async fn bodies_arrive_in_the_chunks_they_were_built_with_or_split() {
318        let body = |response: ReplayResponse| async {
319            let transport = ReplayTransport::new([response]);
320            let response = transport.send(HttpRequest::get("https://x")).await.unwrap();
321            let chunks: Vec<Bytes> = response.into_body().map(Result::unwrap).collect().await;
322            chunks
323        };
324        assert_eq!(
325            body(ReplayResponse::chunks(200, ["ab", "", "cde"])).await,
326            ["ab", "", "cde"]
327        );
328        assert_eq!(
329            body(ReplayResponse::chunks(200, ["ab", "", "cde"]).split_every(2)).await,
330            ["ab", "cd", "e"]
331        );
332        assert_eq!(
333            body(ReplayResponse::new(200, "héllo").split_every(1))
334                .await
335                .len(),
336            "héllo".len()
337        );
338        assert!(
339            body(ReplayResponse::new(200, "").split_every(3))
340                .await
341                .is_empty()
342        );
343    }
344
345    #[test]
346    fn json_takes_the_body_as_written_and_json_value_serializes() {
347        let written = ReplayResponse::json(400, "{not json");
348        assert_eq!(written.chunks, ["{not json"]);
349        let value = ReplayResponse::json_value(200, &serde_json::json!({"ok": true}));
350        assert_eq!(value.chunks, [r#"{"ok":true}"#]);
351        assert_eq!(value.headers.get("content-type"), Some("application/json"));
352    }
353
354    #[tokio::test]
355    async fn stalled_bodies_never_end() {
356        let transport = ReplayTransport::new([ReplayResponse::sse("data: a\n\n").then_stall()]);
357        let mut body = transport
358            .send(HttpRequest::get("https://x"))
359            .await
360            .unwrap()
361            .into_body();
362        assert!(body.next().await.is_some());
363        let next = futures_util::future::select(
364            body.next(),
365            futures_timer::Delay::new(std::time::Duration::from_millis(20)),
366        )
367        .await;
368        assert!(matches!(next, futures_util::future::Either::Right(_)));
369    }
370}
allama_http/replay.rs

allama_http/
replay.rs
