rig_core/test_utils/
http.rs

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

rig_core/test_utils/
http.rs
