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;
7pub mod observations;
8#[cfg(test)]
9pub(crate) mod provider_extensions;
10#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
11pub mod raw_parity;
12mod relay;
13mod streaming;
14#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
15pub mod streaming_conformance;
16#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
17mod streaming_conformance_suite;
18mod trace_capture;
19mod tracing_isolation;
20
21pub use completion::{
22    MOCK_API, MOCK_MODEL, MockCompletionModel, MockError, MockRuntime, MockScript, MockTurn,
23    refuse_options,
24};
25pub use embeddings::{MockEmbeddingModel, MockEmbeddings};
26pub use relay::MockRelay;
27pub use rig_http::test_utils::{
28    CapturedHttpRequest, HttpErrorStreamingClient, MockHttpResponse, MockStreamingClient,
29    NonSuccessStreamingClient, RecordingHttpClient, SequencedHttpClient,
30    SequencedStreamingHttpClient,
31};
32pub use streaming::{
33    MOCK_PROVIDER, MockDecoder, MockDocument, MockFrame, MockStreamEvent, mock_final,
34    mock_final_with_total_tokens,
35};
36pub use trace_capture::{CapturedEvent, CapturedSpan, TraceCapture};
37pub use tracing_isolation::{
38    scoped_tracing_subscriber_guard, scoped_tracing_subscriber_guard_blocking,
39};
40
41/// The JSON document an encoded request sends.
42///
43/// # Panics
44///
45/// When the body is multipart or is not JSON.
46#[cfg(test)]
47pub(crate) fn json_body(request: &http::Request<crate::wire::Body>) -> serde_json::Value {
48    let crate::wire::Body::Bytes(bytes) = request.body() else {
49        panic!("the request body is multipart, not JSON");
50    };
51    serde_json::from_slice(bytes).expect("the request body is JSON")
52}
53
54/// [`decode_reply`](crate::wire::decode_reply), with the reply's whole
55/// document as a transport reports it (`Null` for none).
56#[cfg(test)]
57pub(crate) use crate::driver::decode_reply;
58
59/// The documents a prepared request sent, as `(id, text)`: `prepare` folds
60/// a request's documents into its first user message as `<file id: ...>`
61/// text documents.
62pub fn sent_documents(request: &crate::completion::CompletionRequest) -> Vec<(String, String)> {
63    request
64        .chat_history
65        .iter()
66        .flat_map(|message| match &**message {
67            crate::message::Message::User { content } => content.iter().collect::<Vec<_>>(),
68            crate::message::Message::System { .. } | crate::message::Message::Assistant(_) => {
69                Vec::new()
70            }
71        })
72        .filter_map(|part| match part {
73            crate::message::UserContent::Document(document) => match &document.data {
74                crate::message::DocumentData::Text(text) => {
75                    let body = text.strip_prefix("<file id: ")?;
76                    let (id, rest) = body.split_once(">\n")?;
77                    let text = rest.strip_suffix("\n</file>\n")?;
78                    Some((id.to_owned(), text.to_owned()))
79                }
80                _ => None,
81            },
82            _ => None,
83        })
84        .collect()
85}
rig_core/test_utils/mod.rs

rig_core/test_utils/
mod.rs
