rig_core/test_utils/
mod.rs

1//! Test utilities for deterministic completion-model tests.
2
3mod completion;
4mod embeddings;
5pub mod history;
6pub mod history_conformance;
7mod memory;
8pub mod observations;
9#[cfg(test)]
10pub(crate) mod provider_extensions;
11#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
12pub mod raw_parity;
13mod relay;
14mod streaming;
15#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
16pub mod streaming_conformance;
17#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
18mod streaming_conformance_suite;
19mod trace_capture;
20mod tracing_isolation;
21
22pub use completion::{
23    MOCK_API, MOCK_MODEL, MockCompletionModel, MockError, MockRuntime, MockScript, MockTurn,
24    refuse_options,
25};
26pub use embeddings::{MockEmbeddingModel, MockEmbeddings, MockMultiTextDocument, MockTextDocument};
27pub use memory::{AppendFailingMemory, CountingMemory, FailingMemory};
28pub use relay::MockRelay;
29pub use rig_http::test_utils::{
30    CapturedHttpRequest, HttpErrorStreamingClient, MockHttpResponse, MockStreamingClient,
31    NonSuccessStreamingClient, RecordingHttpClient, SequencedHttpClient,
32    SequencedStreamingHttpClient,
33};
34pub use streaming::{
35    MOCK_PROVIDER, MockDecoder, MockDocument, MockFrame, MockStreamEvent, mock_final,
36    mock_final_with_total_tokens,
37};
38pub use trace_capture::{CapturedEvent, CapturedSpan, TraceCapture};
39pub use tracing_isolation::{
40    scoped_tracing_subscriber_guard, scoped_tracing_subscriber_guard_blocking,
41};
42
43/// The JSON document an encoded request sends.
44///
45/// # Panics
46///
47/// When the body is multipart or is not JSON.
48#[cfg(test)]
49pub(crate) fn json_body(request: &http::Request<crate::wire::Body>) -> serde_json::Value {
50    let crate::wire::Body::Bytes(bytes) = request.body() else {
51        panic!("the request body is multipart, not JSON");
52    };
53    serde_json::from_slice(bytes).expect("the request body is JSON")
54}
55
56/// Decode one reply of `wire` to `request` from frames already in hand,
57/// folded as `mode` folds it: the one decoder and fold a live call runs.
58#[cfg(test)]
59pub(crate) fn decode_reply<W: crate::wire::Wire>(
60    wire: &W,
61    request: &crate::wire::Request<W>,
62    mode: crate::wire::Mode,
63    frames: impl IntoIterator<Item = W::Frame>,
64    raw: serde_json::Value,
65) -> Result<crate::wire::Response<W>, crate::error::ProviderError> {
66    let shared = std::sync::Mutex::new(crate::wire::Shared::new(fold_for(request, wire, mode)));
67    let fed = crate::driver::feed(
68        &mut wire.decoder(),
69        Some(wire.reassembler()),
70        &shared,
71        frames,
72    );
73    crate::driver::settle(
74        shared,
75        fed,
76        crate::wire::Reply {
77            provider: wire.describe().name.to_owned(),
78            raw,
79            provider_request_id: None,
80        },
81    )
82    .outcome
83}
84
85/// The fold a call to `wire` in `mode` opens for `request`, for tests that
86/// drive a decoder by hand.
87#[cfg(test)]
88pub(crate) fn fold_for<W: crate::wire::Wire>(
89    request: &crate::wire::Request<W>,
90    wire: &W,
91    mode: crate::wire::Mode,
92) -> <W::Op as crate::wire::Operation>::Fold {
93    <W::Op as crate::wire::Operation>::fold(
94        request,
95        &mut crate::wire::Call::new(&wire.describe(), mode),
96    )
97}
98
99/// The documents a prepared request sent, as `(id, text)`: `prepare` folds
100/// a request's documents into its first user message as `<file id: ...>`
101/// text documents.
102pub fn sent_documents(request: &crate::completion::CompletionRequest) -> Vec<(String, String)> {
103    request
104        .chat_history
105        .iter()
106        .flat_map(|message| match message {
107            crate::message::Message::User { content } => content.iter().collect::<Vec<_>>(),
108            crate::message::Message::System { .. } | crate::message::Message::Assistant(_) => {
109                Vec::new()
110            }
111        })
112        .filter_map(|part| match part {
113            crate::message::UserContent::Document(document) => match &document.data {
114                crate::message::DocumentSourceKind::String(text) => {
115                    let body = text.strip_prefix("<file id: ")?;
116                    let (id, rest) = body.split_once(">\n")?;
117                    let text = rest.strip_suffix("\n</file>\n")?;
118                    Some((id.to_owned(), text.to_owned()))
119                }
120                _ => None,
121            },
122            _ => None,
123        })
124        .collect()
125}
rig_core/test_utils/mod.rs

rig_core/test_utils/
mod.rs
