boxmaker/docs/plans/M5b/files/crates/loopd/tests/selftest.rs

//! Tests for the startup self-test, against the fake server. Do not edit.

mod support;

use loopd::llama::Client;
use loopd::selftest::{SelfTestError, run};
use support::{FakeServer, Reply, tokenizer};

const CHAT: &str = "/v1/chat/completions";
const TOKENIZE: &str = "/tokenize";

/// A change to the config before the self-test runs.
type Edit = fn(&mut loopd::config::Config);

/// The recorded props and the fake tokenizer, and no completion yet.
fn props_and_tags() -> FakeServer {
    let server = FakeServer::start();
    server.route("/props", vec![Reply::fixture("props")]);
    server.respond(TOKENIZE, tokenizer);
    server
}

/// A server that passes: the recorded props, the fake tokenizer, then the three recorded
/// completions in order.
fn healthy() -> FakeServer {
    let server = props_and_tags();
    server.route(
        CHAT,
        vec![
            Reply::fixture("tool_call"),
            Reply::fixture("turn1"),
            Reply::fixture("turn2"),
        ],
    );
    server
}

fn run_with(
    server: &FakeServer,
    edit: impl FnOnce(&mut loopd::config::Config),
) -> (Result<(), SelfTestError>, Vec<String>) {
    let mut cfg = support::test_config(&server.socket);
    edit(&mut cfg);
    let mut steps = Vec::new();
    let result = run(&Client::new(cfg), &mut |s| steps.push(s.to_string()));
    (result, steps)
}

#[test]
fn a_healthy_server_passes_all_four_checks_in_order() {
    let server = healthy();
    let (result, steps) = run_with(&server, |_| {});
    assert!(result.is_ok(), "{result:?}");
    assert_eq!(
        steps,
        [
            "server matches config",
            "tags tokenize as recorded",
            "tool call round trip",
            "turn 2 cache hit"
        ]
    );

    let chats = server.requests_to(CHAT);
    assert_eq!(chats.len(), 3);
    // The tool-call check offers exactly one tool; the cache check offers none.
    assert_eq!(chats[0].json()["tools"].as_array().map(Vec::len), Some(1));
    assert!(chats[1].json().get("tools").is_none());
    // All three run on the main slot.
    assert!(chats.iter().all(|c| c.json()["id_slot"] == 0));
    // Turn 2 is turn 1 plus the assistant's answer plus a new question: a strict extension.
    let turn1 = chats[1].json()["messages"].as_array().unwrap().clone();
    let turn2 = chats[2].json()["messages"].as_array().unwrap().clone();
    assert_eq!(turn2.len(), turn1.len() + 2);
    assert_eq!(turn2[..turn1.len()], turn1[..]);
    assert_eq!(
        turn2[turn1.len()],
        serde_json::json!({"role": "assistant", "content": "Blue"})
    );
}

#[test]
fn the_main_slot_comes_from_the_config() {
    let server = healthy();
    let (result, _) = run_with(&server, |c| c.slots.main = 1);
    assert!(result.is_ok(), "{result:?}");
    assert!(
        server
            .requests_to(CHAT)
            .iter()
            .all(|c| c.json()["id_slot"] == 1)
    );
}

#[test]
fn each_expected_value_is_checked() {
    let zero = bxcore::Hash32::ZERO;
    let cases: [(&str, Edit); 3] = [
        ("chat template sha256", |c| {
            c.expect.template_sha256 = bxcore::Hash32::ZERO
        }),
        ("context per slot", |c| c.expect.n_ctx = 4096),
        ("slot count", |c| c.expect.slots = 3),
    ];
    for (want_what, edit) in cases {
        let server = healthy();
        let (result, steps) = run_with(&server, edit);
        match result {
            Err(SelfTestError::Mismatch {
                what,
                expected,
                got,
            }) => {
                assert_eq!(what, want_what);
                assert_ne!(expected, got);
                if what == "chat template sha256" {
                    assert_eq!(expected, zero.to_hex());
                }
            }
            other => panic!("{want_what}: {other:?}"),
        }
        assert_eq!(
            steps.len(),
            1,
            "a server that is not the expected one is not sent any prompt"
        );
        assert!(server.requests_to(CHAT).is_empty());
        assert!(server.requests_to(TOKENIZE).is_empty());
    }
}

#[test]
fn a_completion_without_the_tool_call_fails_the_tool_call_check() {
    let server = props_and_tags();
    server.route(CHAT, vec![Reply::fixture("plain")]);
    let (result, steps) = run_with(&server, |_| {});
    assert!(
        matches!(result, Err(SelfTestError::ToolCall(_))),
        "{result:?}"
    );
    assert_eq!(steps.len(), 3);
}

#[test]
fn a_cold_second_turn_fails_the_cache_check() {
    // turn1 twice: the "second turn" reuses 15 tokens where 46 were left.
    let server = props_and_tags();
    server.route(
        CHAT,
        vec![
            Reply::fixture("tool_call"),
            Reply::fixture("progress"),
            Reply::fixture("turn1"),
        ],
    );
    let (result, steps) = run_with(&server, |_| {});
    match result {
        Err(SelfTestError::CacheMiss { expected, got }) => assert_eq!((expected, got), (7052, 15)),
        other => panic!("{other:?}"),
    }
    assert_eq!(steps.len(), 4);
}

#[test]
fn a_server_that_is_down_is_an_inference_error_after_retries() {
    let server = props_and_tags();
    server.route(CHAT, vec![Reply::fixture("tool_call").cut_after(500)]);
    let (result, _) = run_with(&server, |c| c.limits.retry_attempts = 2);
    assert!(matches!(result, Err(SelfTestError::Infer(_))), "{result:?}");
    assert_eq!(
        server.requests_to(CHAT).len(),
        2,
        "the self-test retries like everything else"
    );
    let e: Box<dyn std::error::Error> = Box::new(result.unwrap_err());
    assert!(!e.to_string().is_empty());
}