rig_core/test_utils/
http.rs

1//! HTTP client doubles for provider tests.
2
3use std::{
4    future::{self, Future},
5    sync::{Arc, Mutex, MutexGuard},
6};
7
8use bytes::Bytes;
9
10use crate::{
11    http_client::{
12        self, HttpClientExt, LazyBody, MultipartForm, Request, Response, StreamingResponse,
13    },
14    wasm_compat::WasmCompatSend,
15};
16
17/// Request data captured by [`RecordingHttpClient`].
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct CapturedHttpRequest {
20    /// Request URI.
21    pub uri: String,
22    /// Request headers.
23    pub headers: http::HeaderMap,
24    /// Request body bytes.
25    pub body: Bytes,
26}
27
28/// Response scripted for [`RecordingHttpClient`].
29#[derive(Clone, Debug)]
30pub enum MockHttpResponse {
31    /// Return this body with a successful HTTP status.
32    Success(Bytes),
33    /// Return a status-code error with the given body text.
34    Error(http::StatusCode, String),
35    /// Return an HTTP response with the given (typically non-success) status
36    /// and body, instead of a transport-level error.
37    ErrorResponse(http::StatusCode, Bytes),
38}
39
40impl MockHttpResponse {
41    /// Create a successful response from bytes.
42    pub fn success(body: impl Into<Bytes>) -> Self {
43        Self::Success(body.into())
44    }
45
46    /// Create an error response with a status code and message.
47    pub fn error(status: http::StatusCode, message: impl Into<String>) -> Self {
48        Self::Error(status, message.into())
49    }
50}
51
52impl Default for MockHttpResponse {
53    fn default() -> Self {
54        Self::Success(Bytes::new())
55    }
56}
57
58/// An [`HttpClientExt`] implementation that records unary requests and returns
59/// a fixed response.
60#[derive(Clone, Debug, Default)]
61pub struct RecordingHttpClient {
62    requests: Arc<Mutex<Vec<CapturedHttpRequest>>>,
63    response: Arc<Mutex<MockHttpResponse>>,
64}
65
66impl RecordingHttpClient {
67    /// Create a client that returns `response_body` for unary requests.
68    pub fn new(response_body: impl Into<Bytes>) -> Self {
69        Self {
70            requests: Arc::new(Mutex::new(Vec::new())),
71            response: Arc::new(Mutex::new(MockHttpResponse::success(response_body))),
72        }
73    }
74
75    /// Create a client that returns an HTTP status error for unary requests.
76    pub fn with_error(status: http::StatusCode, message: impl Into<String>) -> Self {
77        Self {
78            requests: Arc::new(Mutex::new(Vec::new())),
79            response: Arc::new(Mutex::new(MockHttpResponse::error(status, message))),
80        }
81    }
82
83    /// Create a client that returns a non-success HTTP response (status and body)
84    /// for unary requests, instead of a transport-level error.
85    pub fn with_error_response(status: http::StatusCode, body: impl Into<Bytes>) -> Self {
86        Self {
87            requests: Arc::new(Mutex::new(Vec::new())),
88            response: Arc::new(Mutex::new(MockHttpResponse::ErrorResponse(
89                status,
90                body.into(),
91            ))),
92        }
93    }
94
95    /// Return the requests captured so far.
96    pub fn requests(&self) -> Vec<CapturedHttpRequest> {
97        self.requests_guard().clone()
98    }
99
100    /// Replace the scripted unary response.
101    pub fn set_response(&self, response: MockHttpResponse) {
102        *self.response_guard() = response;
103    }
104
105    fn requests_guard(&self) -> MutexGuard<'_, Vec<CapturedHttpRequest>> {
106        match self.requests.lock() {
107            Ok(guard) => guard,
108            Err(poisoned) => poisoned.into_inner(),
109        }
110    }
111
112    fn response_guard(&self) -> MutexGuard<'_, MockHttpResponse> {
113        match self.response.lock() {
114            Ok(guard) => guard,
115            Err(poisoned) => poisoned.into_inner(),
116        }
117    }
118
119    fn record_request(&self, uri: String, headers: http::HeaderMap, body: Bytes) {
120        self.requests_guard()
121            .push(CapturedHttpRequest { uri, headers, body });
122    }
123
124    fn build_unary_response<U>(
125        response: MockHttpResponse,
126    ) -> http_client::Result<Response<LazyBody<U>>>
127    where
128        U: From<Bytes> + WasmCompatSend + 'static,
129    {
130        let (status, response_body) = match response {
131            MockHttpResponse::Success(response_body) => (http::StatusCode::OK, response_body),
132            MockHttpResponse::Error(status, message) => {
133                return Err(http_client::Error::InvalidStatusCodeWithMessage(
134                    status, message,
135                ));
136            }
137            MockHttpResponse::ErrorResponse(status, response_body) => (status, response_body),
138        };
139        let body: LazyBody<U> = Box::pin(async move { Ok(U::from(response_body)) });
140        Response::builder()
141            .status(status)
142            .body(body)
143            .map_err(http_client::Error::Protocol)
144    }
145}
146
147impl HttpClientExt for RecordingHttpClient {
148    fn send<T, U>(
149        &self,
150        req: Request<T>,
151    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
152    where
153        T: Into<Bytes> + WasmCompatSend,
154        U: From<Bytes> + WasmCompatSend + 'static,
155    {
156        let response = self.response_guard().clone();
157        let (parts, body) = req.into_parts();
158        self.record_request(parts.uri.to_string(), parts.headers, body.into());
159
160        async move { Self::build_unary_response(response) }
161    }
162
163    fn send_multipart<U>(
164        &self,
165        req: Request<MultipartForm>,
166    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
167    where
168        U: From<Bytes> + WasmCompatSend + 'static,
169    {
170        let response = self.response_guard().clone();
171        let (parts, _body) = req.into_parts();
172        self.record_request(parts.uri.to_string(), parts.headers, Bytes::new());
173
174        async move { Self::build_unary_response(response) }
175    }
176
177    fn send_streaming<T>(
178        &self,
179        _req: Request<T>,
180    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + WasmCompatSend
181    where
182        T: Into<Bytes> + WasmCompatSend,
183    {
184        future::ready(Err(http_client::Error::InvalidStatusCode(
185            http::StatusCode::NOT_IMPLEMENTED,
186        )))
187    }
188}
189
190/// A mock HTTP client that returns pre-built SSE bytes from `send_streaming`.
191///
192/// `send` and `send_multipart` always return `NOT_IMPLEMENTED`.
193#[derive(Clone, Debug, Default)]
194pub struct MockStreamingClient {
195    /// Bytes returned as a single streaming response chunk.
196    pub sse_bytes: Bytes,
197}
198
199impl HttpClientExt for MockStreamingClient {
200    fn send<T, U>(
201        &self,
202        _req: Request<T>,
203    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
204    where
205        T: Into<Bytes> + WasmCompatSend,
206        U: From<Bytes> + WasmCompatSend + 'static,
207    {
208        future::ready(Err(http_client::Error::InvalidStatusCode(
209            http::StatusCode::NOT_IMPLEMENTED,
210        )))
211    }
212
213    fn send_multipart<U>(
214        &self,
215        _req: Request<MultipartForm>,
216    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
217    where
218        U: From<Bytes> + WasmCompatSend + 'static,
219    {
220        future::ready(Err(http_client::Error::InvalidStatusCode(
221            http::StatusCode::NOT_IMPLEMENTED,
222        )))
223    }
224
225    fn send_streaming<T>(
226        &self,
227        _req: Request<T>,
228    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + WasmCompatSend
229    where
230        T: Into<Bytes> + WasmCompatSend,
231    {
232        let sse_bytes = self.sse_bytes.clone();
233        async move {
234            let byte_stream =
235                futures::stream::iter(vec![Ok::<Bytes, http_client::Error>(sse_bytes)]);
236            let boxed_stream: http_client::sse::BoxedStream = Box::pin(byte_stream);
237
238            Response::builder()
239                .status(http::StatusCode::OK)
240                .header(http::header::CONTENT_TYPE, "text/event-stream")
241                .body(boxed_stream)
242                .map_err(http_client::Error::Protocol)
243        }
244    }
245}
246
247/// An [`HttpClientExt`] implementation whose `send_streaming` fails immediately
248/// with a non-success HTTP status and response body.
249#[derive(Debug, Clone)]
250pub struct HttpErrorStreamingClient {
251    pub status: http::StatusCode,
252    pub body: String,
253}
254
255impl HttpErrorStreamingClient {
256    /// Create a streaming client that fails `send_streaming` with the given status and body.
257    pub fn new(status: http::StatusCode, body: impl Into<String>) -> Self {
258        Self {
259            status,
260            body: body.into(),
261        }
262    }
263}
264
265impl Default for HttpErrorStreamingClient {
266    /// The completion-model client bound requires `H: Default`; this lets the
267    /// streaming error client back a real model in tests.
268    fn default() -> Self {
269        Self::new(http::StatusCode::INTERNAL_SERVER_ERROR, String::new())
270    }
271}
272
273impl HttpClientExt for HttpErrorStreamingClient {
274    fn send<T, U>(
275        &self,
276        _req: Request<T>,
277    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
278    where
279        T: Into<Bytes> + WasmCompatSend,
280        U: From<Bytes> + WasmCompatSend + 'static,
281    {
282        future::ready(Err(http_client::Error::InvalidStatusCode(
283            http::StatusCode::NOT_IMPLEMENTED,
284        )))
285    }
286
287    fn send_multipart<U>(
288        &self,
289        _req: Request<MultipartForm>,
290    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
291    where
292        U: From<Bytes> + WasmCompatSend + 'static,
293    {
294        future::ready(Err(http_client::Error::InvalidStatusCode(
295            http::StatusCode::NOT_IMPLEMENTED,
296        )))
297    }
298
299    fn send_streaming<T>(
300        &self,
301        _req: Request<T>,
302    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + WasmCompatSend
303    where
304        T: Into<Bytes> + WasmCompatSend,
305    {
306        let status = self.status;
307        let body = self.body.clone();
308        async move {
309            Err(http_client::Error::InvalidStatusCodeWithMessage(
310                status, body,
311            ))
312        }
313    }
314}
315
316/// An [`HttpClientExt`] implementation that returns one scripted stream of byte
317/// chunks from `send_streaming`.
318#[derive(Debug, Clone, Default)]
319pub struct SequencedStreamingHttpClient {
320    chunks: Arc<Mutex<Option<Vec<http_client::Result<Bytes>>>>>,
321}
322
323impl SequencedStreamingHttpClient {
324    /// Create a streaming client from the chunks it should yield.
325    pub fn new(chunks: Vec<http_client::Result<Bytes>>) -> Self {
326        Self {
327            chunks: Arc::new(Mutex::new(Some(chunks))),
328        }
329    }
330}
331
332impl HttpClientExt for SequencedStreamingHttpClient {
333    fn send<T, U>(
334        &self,
335        _req: Request<T>,
336    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
337    where
338        T: Into<Bytes> + WasmCompatSend,
339        U: From<Bytes> + WasmCompatSend + 'static,
340    {
341        future::ready(Err(http_client::Error::InvalidStatusCode(
342            http::StatusCode::NOT_IMPLEMENTED,
343        )))
344    }
345
346    fn send_multipart<U>(
347        &self,
348        _req: Request<MultipartForm>,
349    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + WasmCompatSend + 'static
350    where
351        U: From<Bytes> + WasmCompatSend + 'static,
352    {
353        future::ready(Err(http_client::Error::InvalidStatusCode(
354            http::StatusCode::NOT_IMPLEMENTED,
355        )))
356    }
357
358    fn send_streaming<T>(
359        &self,
360        _req: Request<T>,
361    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + WasmCompatSend
362    where
363        T: Into<Bytes> + WasmCompatSend,
364    {
365        let chunks = match self.chunks.lock() {
366            Ok(mut guard) => guard.take(),
367            Err(poisoned) => poisoned.into_inner().take(),
368        };
369
370        async move {
371            let Some(chunks) = chunks else {
372                return Err(http_client::Error::InvalidStatusCodeWithMessage(
373                    http::StatusCode::INTERNAL_SERVER_ERROR,
374                    "streaming chunks should only be consumed once".to_string(),
375                ));
376            };
377
378            let byte_stream = futures::stream::iter(chunks);
379            let boxed_stream: http_client::sse::BoxedStream = Box::pin(byte_stream);
380
381            Response::builder()
382                .status(http::StatusCode::OK)
383                .header(http::header::CONTENT_TYPE, "text/event-stream")
384                .body(boxed_stream)
385                .map_err(http_client::Error::Protocol)
386        }
387    }
388}
rig_core/test_utils/http.rs

rig_core/test_utils/
http.rs
