Skip to main content

rig_core/test_utils/
mod.rs

1//! Test utilities for deterministic completion-model tests.
2
3mod completion;
4mod embeddings;
5mod memory;
6pub mod observations;
7mod relay;
8mod streaming;
9#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
10pub mod streaming_conformance;
11#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
12mod streaming_conformance_suite;
13mod tracing_isolation;
14
15pub use completion::{MockCompletionModel, MockError, MockRuntime, MockScript, MockTurn};
16pub use embeddings::{MockEmbeddingModel, MockEmbeddings, MockMultiTextDocument, MockTextDocument};
17pub use memory::{AppendFailingMemory, CountingMemory, FailingMemory};
18pub use relay::MockRelay;
19pub use rig_http::test_utils::{
20    CapturedHttpRequest, HttpErrorStreamingClient, MockHttpResponse, MockStreamingClient,
21    NonSuccessStreamingClient, RecordingHttpClient, SequencedHttpClient,
22    SequencedStreamingHttpClient,
23};
24pub use streaming::{
25    MOCK_PROVIDER, MockDecoder, MockFrame, MockStreamEvent, mock_final,
26    mock_final_with_total_tokens,
27};
28pub use tracing_isolation::{
29    scoped_tracing_subscriber_guard, scoped_tracing_subscriber_guard_blocking,
30};
31
32/// Decode one reply of `wire` to `request` from frames already in hand,
33/// folded as `mode` folds it: the one decoder and fold a live call runs.
34#[cfg(test)]
35pub(crate) fn decode_reply<W: crate::wire::Wire>(
36    wire: &W,
37    request: &crate::wire::Request<W>,
38    mode: crate::wire::Mode,
39    frames: impl IntoIterator<Item = W::Frame>,
40    raw: serde_json::Value,
41) -> Result<crate::wire::Response<W>, crate::error::ProviderError> {
42    crate::driver::decode_frames(
43        wire,
44        fold_for(request, wire, mode),
45        frames,
46        crate::wire::Reply {
47            provider: wire.describe().name.to_owned(),
48            raw,
49            provider_request_id: None,
50        },
51    )
52}
53
54/// The fold a call to `wire` in `mode` opens for `request`, for tests that
55/// drive a decoder by hand.
56#[cfg(test)]
57pub(crate) fn fold_for<W: crate::wire::Wire>(
58    request: &crate::wire::Request<W>,
59    wire: &W,
60    mode: crate::wire::Mode,
61) -> <W::Op as crate::wire::Operation>::Fold {
62    <W::Op as crate::wire::Operation>::fold(
63        request,
64        &mut crate::wire::Call::new(&wire.describe(), mode),
65    )
66}