typesafe_rs_mock/
lib.rs

1//! In-process HTTP mock for [`typesafe-rs`](https://docs.rs/typesafe-rs).
2//!
3//! Binds `127.0.0.1:0`, scripts `POST /v1/systemone` and `GET /v1/models`, and
4//! records a request journal. Tests use a real `Client` over loopback.
5
6#![forbid(unsafe_code)]
7#![deny(missing_docs)]
8#![warn(missing_debug_implementations)]
9
10use std::collections::HashMap;
11use std::fmt;
12use std::sync::{Arc, Mutex};
13use std::time::{Duration, Instant};
14
15use axum::Router;
16use axum::body::{Body, to_bytes};
17use axum::extract::State;
18use axum::http::{HeaderMap as AxumHeaderMap, HeaderValue, Request, StatusCode};
19use axum::response::{IntoResponse, Response};
20use axum::routing::{get, post};
21use serde_json::{Value, json};
22use tokio::sync::oneshot;
23use url::Url;
24
25/// Running mock API, bound to a random localhost port.
26pub struct MockServer {
27    url: Url,
28    inner: Arc<Mutex<Inner>>,
29    shutdown: Option<oneshot::Sender<()>>,
30    join: Option<tokio::task::JoinHandle<()>>,
31}
32
33struct Inner {
34    stubs: Vec<Stub>,
35    journal: Vec<RecordedRequest>,
36    request_counter: u64,
37}
38
39#[derive(Clone)]
40struct Stub {
41    endpoint: Endpoint,
42    matcher: Matcher,
43    status: u16,
44    headers: Vec<(String, String)>,
45    body: Value,
46    remaining: Option<u32>,
47    delay: Option<Duration>,
48}
49
50#[derive(Clone, Copy, PartialEq, Eq)]
51enum Endpoint {
52    SystemOne,
53    Models,
54}
55
56#[derive(Clone)]
57enum Matcher {
58    Any,
59    QuestionKey(String),
60}
61
62/// One observed HTTP request.
63#[derive(Clone, Debug)]
64pub struct RecordedRequest {
65    /// HTTP method, uppercase.
66    pub method: String,
67    /// Request path, e.g. `/v1/systemone`.
68    pub path: String,
69    /// Headers keyed by lowercase name; last value wins.
70    pub headers: HashMap<String, String>,
71    /// Parsed JSON body, when present and valid.
72    pub body: Option<Value>,
73    /// Local time the mock accepted the request.
74    pub received_at: Instant,
75}
76
77impl RecordedRequest {
78    /// Header value by lowercase name.
79    #[must_use]
80    pub fn header(&self, name: &str) -> Option<&str> {
81        self.headers
82            .get(&name.to_ascii_lowercase())
83            .map(String::as_str)
84    }
85}
86
87impl fmt::Debug for MockServer {
88    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
89        f.debug_struct("MockServer")
90            .field("url", &self.url)
91            .finish_non_exhaustive()
92    }
93}
94
95impl fmt::Debug for StubBuilder {
96    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97        f.debug_struct("StubBuilder")
98            .field("mounted", &self.mounted)
99            .finish_non_exhaustive()
100    }
101}
102
103impl MockServer {
104    /// Bind `127.0.0.1:0` and serve until dropped.
105    pub async fn start() -> Self {
106        let inner = Arc::new(Mutex::new(Inner {
107            stubs: Vec::new(),
108            journal: Vec::new(),
109            request_counter: 0,
110        }));
111        let app = Router::new()
112            .route("/v1/systemone", post(system_one))
113            .route("/v1/models", get(models))
114            .fallback(fallback)
115            .with_state(inner.clone());
116
117        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
118            .await
119            .expect("bind mock listener");
120        let addr = listener.local_addr().expect("local_addr");
121        let url = Url::parse(&format!("http://{addr}")).expect("url");
122        let (tx, rx) = oneshot::channel();
123        let join = tokio::spawn(async move {
124            let _ = axum::serve(listener, app)
125                .with_graceful_shutdown(async {
126                    let _ = rx.await;
127                })
128                .await;
129        });
130
131        Self {
132            url,
133            inner,
134            shutdown: Some(tx),
135            join: Some(join),
136        }
137    }
138
139    /// Base URL with no trailing slash, e.g. `http://127.0.0.1:12345`.
140    #[must_use]
141    pub fn url(&self) -> Url {
142        self.url.clone()
143    }
144
145    /// Script `POST /v1/systemone` responses. The stub is registered when the builder is dropped.
146    pub fn on_system_one(&self) -> StubBuilder {
147        StubBuilder::new(self.inner.clone(), Endpoint::SystemOne)
148    }
149
150    /// Script `GET /v1/models` responses. The stub is registered when the builder is dropped.
151    pub fn on_models(&self) -> StubBuilder {
152        StubBuilder::new(self.inner.clone(), Endpoint::Models)
153    }
154
155    /// Snapshot of received requests, in order.
156    #[must_use]
157    pub fn journal(&self) -> Vec<RecordedRequest> {
158        self.inner
159            .lock()
160            .unwrap_or_else(|e| e.into_inner())
161            .journal
162            .clone()
163    }
164
165    /// Gaps between consecutive journal timestamps.
166    #[must_use]
167    pub fn request_gaps(&self) -> Vec<Duration> {
168        let journal = self.journal();
169        journal
170            .windows(2)
171            .map(|w| w[1].received_at.saturating_duration_since(w[0].received_at))
172            .collect()
173    }
174}
175
176impl Drop for MockServer {
177    fn drop(&mut self) {
178        if let Some(tx) = self.shutdown.take() {
179            let _ = tx.send(());
180        }
181        if let Some(join) = self.join.take() {
182            join.abort();
183        }
184    }
185}
186
187/// Fluent stub registration. Mounted on drop.
188pub struct StubBuilder {
189    inner: Arc<Mutex<Inner>>,
190    stub: Stub,
191    mounted: bool,
192}
193
194impl StubBuilder {
195    fn new(inner: Arc<Mutex<Inner>>, endpoint: Endpoint) -> Self {
196        let body = match endpoint {
197            Endpoint::SystemOne => json!({
198                "model": "jev-latest",
199                "answers": {},
200                "usage": { "input_tokens": 0, "output_tokens": 0 }
201            }),
202            Endpoint::Models => json!({ "models": [] }),
203        };
204        Self {
205            inner,
206            stub: Stub {
207                endpoint,
208                matcher: Matcher::Any,
209                status: 200,
210                headers: Vec::new(),
211                body,
212                remaining: None,
213                delay: None,
214            },
215            mounted: false,
216        }
217    }
218
219    /// Only match systemone requests whose `questions` map contains `key`.
220    pub fn with_question_key(mut self, key: impl Into<String>) -> Self {
221        self.stub.matcher = Matcher::QuestionKey(key.into());
222        self
223    }
224
225    /// JSON body to return.
226    ///
227    /// If `body` looks like an answers map (no top-level `answers` / `error` /
228    /// `models` keys), it is wrapped as a System One response. Use
229    /// [`Self::respond_raw`] to send an exact body.
230    pub fn respond(mut self, body: Value) -> Self {
231        self.stub.body = wrap_systemone_body(self.stub.endpoint, body);
232        self
233    }
234
235    /// JSON body to return, without wrapping.
236    pub fn respond_raw(mut self, body: Value) -> Self {
237        self.stub.body = body;
238        self
239    }
240
241    /// Set the HTTP status (and a small JSON error body when not 2xx).
242    pub fn respond_status(mut self, status: u16) -> Self {
243        self.stub.status = status;
244        if !(200..300).contains(&status) {
245            self.stub.body = json!({ "error": format!("status {status}") });
246        }
247        self
248    }
249
250    /// Add a response header.
251    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
252        self.stub.headers.push((name.into(), value.into()));
253        self
254    }
255
256    /// How many times this stub may match. Omit for unlimited.
257    pub fn times(mut self, n: u32) -> Self {
258        self.stub.remaining = Some(n);
259        self
260    }
261
262    /// Sleep before responding (latency injection).
263    pub fn delay(mut self, delay: Duration) -> Self {
264        self.stub.delay = Some(delay);
265        self
266    }
267
268    fn mount(&mut self) {
269        if self.mounted {
270            return;
271        }
272        self.inner
273            .lock()
274            .unwrap_or_else(|e| e.into_inner())
275            .stubs
276            .push(self.stub.clone());
277        self.mounted = true;
278    }
279}
280
281impl Drop for StubBuilder {
282    fn drop(&mut self) {
283        self.mount();
284    }
285}
286
287fn wrap_systemone_body(endpoint: Endpoint, body: Value) -> Value {
288    if endpoint != Endpoint::SystemOne {
289        return body;
290    }
291    if body.get("answers").is_some() || body.get("error").is_some() || body.get("models").is_some()
292    {
293        return body;
294    }
295    json!({
296        "model": "jev-latest",
297        "answers": body,
298        "usage": { "input_tokens": 0, "output_tokens": 0 }
299    })
300}
301
302/// Build a Noul answer object.
303#[must_use]
304pub fn noul(value: f64) -> Value {
305    json!({ "type": "noul", "noul": value })
306}
307
308/// Build a Choice answer object.
309#[must_use]
310pub fn choice(label: &str, confidence: f64) -> Value {
311    json!({
312        "type": "choice",
313        "choice": label,
314        "probabilities": { label: 1.0 },
315        "confidence": confidence
316    })
317}
318
319/// Build a Score answer object.
320#[must_use]
321pub fn score(value: f64, confidence: f64) -> Value {
322    json!({
323        "type": "score",
324        "score": value,
325        "legend": { "0": "low", "1": "high" },
326        "probabilities": { "0": 1.0 - value.min(1.0), "1": value.min(1.0) },
327        "confidence": confidence
328    })
329}
330
331async fn system_one(State(state): State<Arc<Mutex<Inner>>>, req: Request<Body>) -> Response {
332    handle(state, Endpoint::SystemOne, req).await
333}
334
335async fn models(State(state): State<Arc<Mutex<Inner>>>, req: Request<Body>) -> Response {
336    handle(state, Endpoint::Models, req).await
337}
338
339async fn fallback() -> impl IntoResponse {
340    (
341        StatusCode::NOT_FOUND,
342        axum::Json(json!({ "error": "not found" })),
343    )
344}
345
346async fn handle(state: Arc<Mutex<Inner>>, endpoint: Endpoint, req: Request<Body>) -> Response {
347    let method = req.method().as_str().to_owned();
348    let path = req.uri().path().to_owned();
349    let header_map = req.headers().clone();
350    let (_parts, body) = req.into_parts();
351    let bytes = to_bytes(body, 2 * 1024 * 1024).await.unwrap_or_default();
352    let json_body = serde_json::from_slice::<Value>(&bytes).ok();
353
354    let mut recorded_headers = HashMap::new();
355    for (name, value) in &header_map {
356        if let Ok(v) = value.to_str() {
357            recorded_headers.insert(name.as_str().to_ascii_lowercase(), v.to_owned());
358        }
359    }
360
361    let recorded = RecordedRequest {
362        method,
363        path,
364        headers: recorded_headers,
365        body: json_body.clone(),
366        received_at: Instant::now(),
367    };
368
369    let chosen = {
370        let mut guard = state.lock().unwrap_or_else(|e| e.into_inner());
371        guard.journal.push(recorded);
372        guard.request_counter += 1;
373        let request_id = format!("mock-{}", guard.request_counter);
374        let idx = guard.stubs.iter().position(|stub| {
375            if stub.endpoint != endpoint {
376                return false;
377            }
378            if stub.remaining == Some(0) {
379                return false;
380            }
381            match &stub.matcher {
382                Matcher::Any => true,
383                Matcher::QuestionKey(key) => json_body
384                    .as_ref()
385                    .and_then(|b| b.get("questions"))
386                    .and_then(Value::as_object)
387                    .is_some_and(|q| q.contains_key(key)),
388            }
389        });
390        idx.map(|i| {
391            if let Some(left) = &mut guard.stubs[i].remaining {
392                *left = left.saturating_sub(1);
393            }
394            let mut stub = guard.stubs[i].clone();
395            if !stub
396                .headers
397                .iter()
398                .any(|(n, _)| n.eq_ignore_ascii_case("x-typesafe-request-id"))
399            {
400                stub.headers
401                    .push(("x-typesafe-request-id".to_owned(), request_id));
402            }
403            stub
404        })
405    };
406
407    let Some(stub) = chosen else {
408        return (
409            StatusCode::INTERNAL_SERVER_ERROR,
410            axum::Json(json!({ "error": "no stub matched" })),
411        )
412            .into_response();
413    };
414
415    if let Some(delay) = stub.delay {
416        tokio::time::sleep(delay).await;
417    }
418
419    respond(stub)
420}
421
422fn respond(stub: Stub) -> Response {
423    let status = StatusCode::from_u16(stub.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
424    let body = stub.body.to_string();
425    let mut headers = AxumHeaderMap::new();
426    headers.insert(
427        axum::http::header::CONTENT_TYPE,
428        HeaderValue::from_static("application/json"),
429    );
430    for (name, value) in stub.headers {
431        if let (Ok(n), Ok(v)) = (
432            axum::http::HeaderName::try_from(name),
433            HeaderValue::from_str(&value),
434        ) {
435            headers.insert(n, v);
436        }
437    }
438    (status, headers, body).into_response()
439}
440
441#[cfg(test)]
442mod tests {
443    use super::*;
444
445    #[tokio::test]
446    async fn binds_localhost() {
447        let mock = MockServer::start().await;
448        assert_eq!(mock.url().scheme(), "http");
449        assert_eq!(mock.url().host_str(), Some("127.0.0.1"));
450        assert!(mock.url().port().is_some());
451    }
452}
typesafe_rs_mock/lib.rs

typesafe_rs_mock/
lib.rs
