Skip to main content

rig_http/
test_utils.rs

1//! HTTP client doubles for tests of code that sends through an
2//! [`HttpClientExt`].
3
4use std::{
5    collections::VecDeque,
6    future::{self, Future},
7    sync::{Arc, Mutex, MutexGuard, PoisonError},
8};
9
10use bytes::Bytes;
11
12use crate::{
13    http_client::{
14        self, HttpClientExt, LazyBody, MultipartForm, Request, Response, StreamingResponse,
15    },
16    wasm_compat::WasmCompatSend,
17};
18
19/// The reply a double gives on a surface it does not script: the request
20/// reached a transport that has no answer for it.
21fn not_implemented() -> http_client::Error {
22    http_client::Error::non_success_with_details(
23        http::StatusCode::NOT_IMPLEMENTED,
24        http::HeaderMap::new(),
25        String::new(),
26    )
27}
28
29/// `mutex`'s value, also after a panic poisoned it.
30fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
31    mutex.lock().unwrap_or_else(PoisonError::into_inner)
32}
33
34/// The unary half of a streaming-only double.
35///
36/// Four doubles below script `send_streaming` and nothing else; that they
37/// have no unary surface is one fact about streaming-only doubles, so it is
38/// written here once and each impl states it by name.
39macro_rules! no_unary_surface {
40    () => {
41        fn send<T, U>(
42            &self,
43            _req: Request<T>,
44        ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>>
45        + WasmCompatSend
46        + 'static
47        where
48            T: Into<Bytes> + WasmCompatSend,
49            U: From<Bytes> + WasmCompatSend + 'static,
50        {
51            future::ready(Err(not_implemented()))
52        }
53
54        fn send_multipart<U>(
55            &self,
56            _req: Request<MultipartForm>,
57        ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>>
58        + WasmCompatSend
59        + 'static
60        where
61            U: From<Bytes> + WasmCompatSend + 'static,
62        {
63            future::ready(Err(not_implemented()))
64        }
65    };
66}
67
68/// Request data captured by [`RecordingHttpClient`].
69#[derive(Debug, Clone, PartialEq, Eq)]
70pub struct CapturedHttpRequest {
71    /// Request URI.
72    pub uri: String,
73    /// Request headers.
74    pub headers: http::HeaderMap,
75    /// Request body bytes.
76    pub body: Bytes,
77}
78
79/// Response scripted for [`RecordingHttpClient`].
80#[derive(Clone, Debug)]
81pub enum MockHttpResponse {
82    /// Return this body with a successful HTTP status and no headers.
83    Success(Bytes),
84    /// Return this body with a successful HTTP status and these headers —
85    /// the reply whose `Content-Type` decides how it may be framed.
86    SuccessWithHeaders(Bytes, http::HeaderMap),
87    /// Return an HTTP response with the given (typically non-success) status
88    /// and body, instead of a transport-level error.
89    ErrorResponse(http::StatusCode, Bytes),
90    /// Return a status-code error that preserved the failed response's
91    /// headers, exactly as the bundled reqwest transport does (rig#2210).
92    ErrorWithHeaders(http::StatusCode, String, http::HeaderMap),
93    /// Return an HTTP response with the given (typically non-success) status,
94    /// body, and response headers, instead of a transport-level error.
95    ErrorResponseWithHeaders(http::StatusCode, Bytes, http::HeaderMap),
96}
97
98impl MockHttpResponse {
99    /// Create a successful response from bytes.
100    pub fn success(body: impl Into<Bytes>) -> Self {
101        Self::Success(body.into())
102    }
103
104    /// Create a successful response carrying `content_type`.
105    pub fn success_typed(body: impl Into<Bytes>, content_type: &'static str) -> Self {
106        let mut headers = http::HeaderMap::new();
107        headers.insert(
108            http::header::CONTENT_TYPE,
109            http::HeaderValue::from_static(content_type),
110        );
111        Self::SuccessWithHeaders(body.into(), headers)
112    }
113
114    /// Create a transport-level status error whose reply carried no headers
115    /// of interest: the same shape as [`Self::error_with_headers`] with an
116    /// empty map, since a rejection always comes with its headers.
117    pub fn error(status: http::StatusCode, message: impl Into<String>) -> Self {
118        Self::error_with_headers(status, message, http::HeaderMap::new())
119    }
120
121    /// Create a transport-level status error that preserved the failed
122    /// response's headers, as the bundled reqwest transport does.
123    pub fn error_with_headers(
124        status: http::StatusCode,
125        message: impl Into<String>,
126        headers: http::HeaderMap,
127    ) -> Self {
128        Self::ErrorWithHeaders(status, message.into(), headers)
129    }
130}
131
132impl MockHttpResponse {
133    /// The status, body and headers this reply has, or the transport error
134    /// it is.
135    fn into_parts(self) -> http_client::Result<(http::StatusCode, Bytes, Option<http::HeaderMap>)> {
136        Ok(match self {
137            Self::Success(body) => (http::StatusCode::OK, body, None),
138            Self::SuccessWithHeaders(body, headers) => (http::StatusCode::OK, body, Some(headers)),
139            Self::ErrorWithHeaders(status, body, headers) => {
140                return Err(http_client::Error::InvalidStatusCodeWithDetails {
141                    status,
142                    body,
143                    headers,
144                });
145            }
146            Self::ErrorResponse(status, body) => (status, body, None),
147            Self::ErrorResponseWithHeaders(status, body, headers) => (status, body, Some(headers)),
148        })
149    }
150}
151
152impl Default for MockHttpResponse {
153    fn default() -> Self {
154        Self::Success(Bytes::new())
155    }
156}
157
158/// An [`HttpClientExt`] implementation that records unary requests and returns
159/// a fixed response.
160#[derive(Clone, Debug, Default)]
161pub struct RecordingHttpClient {
162    requests: Arc<Mutex<Vec<CapturedHttpRequest>>>,
163    response: Arc<Mutex<MockHttpResponse>>,
164}
165
166impl RecordingHttpClient {
167    /// A client that answers every unary request with `response`.
168    fn scripted(response: MockHttpResponse) -> Self {
169        Self {
170            requests: Arc::default(),
171            response: Arc::new(Mutex::new(response)),
172        }
173    }
174
175    /// Create a client that returns `response_body` for unary requests.
176    pub fn new(response_body: impl Into<Bytes>) -> Self {
177        Self::scripted(MockHttpResponse::success(response_body))
178    }
179
180    /// Create a client that returns an HTTP status error for unary requests.
181    pub fn with_error(status: http::StatusCode, message: impl Into<String>) -> Self {
182        Self::scripted(MockHttpResponse::error(status, message))
183    }
184
185    /// Create a client that returns a non-success HTTP response (status and body)
186    /// for unary requests, instead of a transport-level error.
187    pub fn with_error_response(status: http::StatusCode, body: impl Into<Bytes>) -> Self {
188        Self::scripted(MockHttpResponse::ErrorResponse(status, body.into()))
189    }
190
191    /// Create a client whose transport reports a non-success status *and*
192    /// preserves the failed response's headers, as the bundled reqwest client
193    /// does (rig#2210).
194    pub fn with_error_headers(
195        status: http::StatusCode,
196        message: impl Into<String>,
197        headers: http::HeaderMap,
198    ) -> Self {
199        Self::scripted(MockHttpResponse::error_with_headers(
200            status, message, headers,
201        ))
202    }
203
204    /// Create a client that hands back a non-success HTTP response carrying
205    /// `headers`, instead of a transport-level error.
206    pub fn with_error_response_headers(
207        status: http::StatusCode,
208        body: impl Into<Bytes>,
209        headers: http::HeaderMap,
210    ) -> Self {
211        Self::scripted(MockHttpResponse::ErrorResponseWithHeaders(
212            status,
213            body.into(),
214            headers,
215        ))
216    }
217
218    /// Return the requests captured so far.
219    pub fn requests(&self) -> Vec<CapturedHttpRequest> {
220        lock(&self.requests).clone()
221    }
222
223    fn record_request(&self, uri: String, headers: http::HeaderMap, body: Bytes) {
224        lock(&self.requests).push(CapturedHttpRequest { uri, headers, body });
225    }
226
227    fn build_unary_response<U>(
228        response: MockHttpResponse,
229    ) -> http_client::Result<Response<LazyBody<U>>>
230    where
231        U: From<Bytes> + WasmCompatSend + 'static,
232    {
233        let (status, response_body, response_headers) = response.into_parts()?;
234        let body: LazyBody<U> = Box::pin(async move { Ok(U::from(response_body)) });
235        let mut builder = Response::builder().status(status);
236        if let Some(headers) = response_headers
237            && let Some(slot) = builder.headers_mut()
238        {
239            *slot = headers;
240        }
241        builder.body(body).map_err(http_client::Error::Protocol)
242    }
243}
244
245impl HttpClientExt for RecordingHttpClient {
246    fn send<T, U>(
247        &self,
248        req: Request<T>,
249    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
250    where
251        T: Into<Bytes> + WasmCompatSend,
252        U: From<Bytes> + WasmCompatSend + 'static,
253    {
254        let response = lock(&self.response).clone();
255        let (parts, body) = req.into_parts();
256        self.record_request(parts.uri.to_string(), parts.headers, body.into());
257
258        async move { Self::build_unary_response(response) }
259    }
260
261    fn send_multipart<U>(
262        &self,
263        req: Request<MultipartForm>,
264    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
265    where
266        U: From<Bytes> + WasmCompatSend + 'static,
267    {
268        let response = lock(&self.response).clone();
269        let (parts, body) = req.into_parts();
270        let (_, body) = body.boundary("recording-http-client").encode();
271        self.record_request(parts.uri.to_string(), parts.headers, body);
272
273        async move { Self::build_unary_response(response) }
274    }
275
276    fn send_streaming<T>(
277        &self,
278        _req: Request<T>,
279    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + WasmCompatSend
280    where
281        T: Into<Bytes> + WasmCompatSend,
282    {
283        future::ready(Err(not_implemented()))
284    }
285}
286
287/// An [`HttpClientExt`] implementation that records requests and returns one
288/// scripted response per request. A streamed request's reply body arrives as
289/// one chunk.
290///
291/// This is useful for testing retry and recovery paths through real provider
292/// request/response conversion without live credentials.
293#[derive(Clone, Debug, Default)]
294pub struct SequencedHttpClient {
295    requests: Arc<Mutex<Vec<CapturedHttpRequest>>>,
296    responses: Arc<Mutex<VecDeque<MockHttpResponse>>>,
297}
298
299impl SequencedHttpClient {
300    /// Create a client that returns the supplied responses in order.
301    pub fn new(responses: impl IntoIterator<Item = MockHttpResponse>) -> Self {
302        Self {
303            requests: Arc::new(Mutex::new(Vec::new())),
304            responses: Arc::new(Mutex::new(responses.into_iter().collect())),
305        }
306    }
307
308    /// Return the requests captured so far.
309    pub fn requests(&self) -> Vec<CapturedHttpRequest> {
310        lock(&self.requests).clone()
311    }
312
313    /// Return the number of scripted responses that have not been consumed.
314    pub fn remaining_responses(&self) -> usize {
315        lock(&self.responses).len()
316    }
317
318    /// Records the request `parts` and `body` and takes the next scripted
319    /// reply.
320    fn answer(&self, parts: http::request::Parts, body: Bytes) -> Option<MockHttpResponse> {
321        lock(&self.requests).push(CapturedHttpRequest {
322            uri: parts.uri.to_string(),
323            headers: parts.headers,
324            body,
325        });
326        lock(&self.responses).pop_front()
327    }
328}
329
330impl HttpClientExt for SequencedHttpClient {
331    fn send<T, U>(
332        &self,
333        req: Request<T>,
334    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
335    where
336        T: Into<Bytes> + WasmCompatSend,
337        U: From<Bytes> + WasmCompatSend + 'static,
338    {
339        let (parts, body) = req.into_parts();
340        let response = self.answer(parts, body.into());
341        async move { RecordingHttpClient::build_unary_response(response.ok_or_else(not_implemented)?) }
342    }
343
344    fn send_multipart<U>(
345        &self,
346        req: Request<MultipartForm>,
347    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
348    where
349        U: From<Bytes> + WasmCompatSend + 'static,
350    {
351        let (parts, _body) = req.into_parts();
352        let response = self.answer(parts, Bytes::new());
353        async move { RecordingHttpClient::build_unary_response(response.ok_or_else(not_implemented)?) }
354    }
355
356    fn send_streaming<T>(
357        &self,
358        req: Request<T>,
359    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + WasmCompatSend
360    where
361        T: Into<Bytes> + WasmCompatSend,
362    {
363        let (parts, body) = req.into_parts();
364        let response = self.answer(parts, body.into());
365        future::ready(
366            response
367                .ok_or_else(not_implemented)
368                .and_then(streaming_response),
369        )
370    }
371}
372
373/// `response` as a streamed reply whose body is one chunk.
374fn streaming_response(response: MockHttpResponse) -> http_client::Result<StreamingResponse> {
375    let (status, body, headers) = response.into_parts()?;
376    let chunks: http_client::BoxedStream = Box::pin(futures::stream::iter([Ok::<
377        Bytes,
378        http_client::Error,
379    >(body)]));
380    let mut response = Response::builder()
381        .status(status)
382        .body(chunks)
383        .map_err(http_client::Error::Protocol)?;
384    if let Some(headers) = headers {
385        *response.headers_mut() = headers;
386    }
387    Ok(response)
388}
389
390/// A mock HTTP client that returns pre-built SSE bytes from `send_streaming`.
391///
392/// `send` and `send_multipart` always return `NOT_IMPLEMENTED`.
393#[derive(Clone, Debug, Default)]
394pub struct MockStreamingClient {
395    /// Bytes returned as a single streaming response chunk.
396    pub sse_bytes: Bytes,
397}
398
399impl HttpClientExt for MockStreamingClient {
400    no_unary_surface!();
401
402    fn send_streaming<T>(
403        &self,
404        _req: Request<T>,
405    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + WasmCompatSend
406    where
407        T: Into<Bytes> + WasmCompatSend,
408    {
409        let sse_bytes = self.sse_bytes.clone();
410        async move {
411            let byte_stream =
412                futures::stream::iter(vec![Ok::<Bytes, http_client::Error>(sse_bytes)]);
413            let boxed_stream: http_client::BoxedStream = Box::pin(byte_stream);
414
415            Response::builder()
416                .status(http::StatusCode::OK)
417                .header(http::header::CONTENT_TYPE, "text/event-stream")
418                .body(boxed_stream)
419                .map_err(http_client::Error::Protocol)
420        }
421    }
422}
423
424/// An [`HttpClientExt`] implementation whose `send_streaming` fails immediately
425/// with a non-success HTTP status and response body.
426#[derive(Debug, Clone)]
427pub struct HttpErrorStreamingClient {
428    pub status: http::StatusCode,
429    pub body: String,
430}
431
432impl HttpErrorStreamingClient {
433    /// Create a streaming client that fails `send_streaming` with the given status and body.
434    pub fn new(status: http::StatusCode, body: impl Into<String>) -> Self {
435        Self {
436            status,
437            body: body.into(),
438        }
439    }
440}
441
442impl Default for HttpErrorStreamingClient {
443    /// The completion-model client bound requires `H: Default`; this lets the
444    /// streaming error client back a real model in tests.
445    fn default() -> Self {
446        Self::new(http::StatusCode::INTERNAL_SERVER_ERROR, String::new())
447    }
448}
449
450impl HttpClientExt for HttpErrorStreamingClient {
451    no_unary_surface!();
452
453    fn send_streaming<T>(
454        &self,
455        _req: Request<T>,
456    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + WasmCompatSend
457    where
458        T: Into<Bytes> + WasmCompatSend,
459    {
460        let status = self.status;
461        let body = self.body.clone();
462        async move {
463            Err(http_client::Error::non_success_with_details(
464                status,
465                http::HeaderMap::new(),
466                body,
467            ))
468        }
469    }
470}
471
472/// An [`HttpClientExt`] whose `send_streaming` hands a non-success reply back
473/// as an `Ok` response — the fourth non-success cell: a custom transport that
474/// does not reject on status, so the driver must.
475#[derive(Debug, Clone)]
476pub struct NonSuccessStreamingClient {
477    pub status: http::StatusCode,
478    pub headers: http::HeaderMap,
479    pub body: Bytes,
480}
481
482impl HttpClientExt for NonSuccessStreamingClient {
483    no_unary_surface!();
484
485    fn send_streaming<T>(
486        &self,
487        _req: Request<T>,
488    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + WasmCompatSend
489    where
490        T: Into<Bytes> + WasmCompatSend,
491    {
492        let status = self.status;
493        let headers = self.headers.clone();
494        let body = self.body.clone();
495        async move {
496            let byte_stream = futures::stream::iter(vec![Ok::<Bytes, http_client::Error>(body)]);
497            let boxed_stream: http_client::BoxedStream = Box::pin(byte_stream);
498            let mut response = Response::builder()
499                .status(status)
500                .body(boxed_stream)
501                .map_err(http_client::Error::Protocol)?;
502            *response.headers_mut() = headers;
503            Ok(response)
504        }
505    }
506}
507
508/// An [`HttpClientExt`] implementation that returns one scripted stream of byte
509/// chunks from `send_streaming`.
510#[derive(Debug, Clone, Default)]
511pub struct SequencedStreamingHttpClient {
512    chunks: Arc<Mutex<Option<Vec<http_client::Result<Bytes>>>>>,
513}
514
515impl SequencedStreamingHttpClient {
516    /// Create a streaming client from the chunks it should yield.
517    pub fn new(chunks: Vec<http_client::Result<Bytes>>) -> Self {
518        Self {
519            chunks: Arc::new(Mutex::new(Some(chunks))),
520        }
521    }
522}
523
524impl HttpClientExt for SequencedStreamingHttpClient {
525    no_unary_surface!();
526
527    fn send_streaming<T>(
528        &self,
529        _req: Request<T>,
530    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + WasmCompatSend
531    where
532        T: Into<Bytes> + WasmCompatSend,
533    {
534        let chunks = lock(&self.chunks).take();
535
536        async move {
537            let Some(chunks) = chunks else {
538                return Err(http_client::Error::non_success_with_details(
539                    http::StatusCode::INTERNAL_SERVER_ERROR,
540                    http::HeaderMap::new(),
541                    "streaming chunks should only be consumed once".to_string(),
542                ));
543            };
544
545            let byte_stream = futures::stream::iter(chunks);
546            let boxed_stream: http_client::BoxedStream = Box::pin(byte_stream);
547
548            Response::builder()
549                .status(http::StatusCode::OK)
550                .header(http::header::CONTENT_TYPE, "text/event-stream")
551                .body(boxed_stream)
552                .map_err(http_client::Error::Protocol)
553        }
554    }
555}