loadmaster_probe/
fixture.rs

1//! Declarative probe fixtures: YAML with request steps, a pass condition and ordered
2//! failure classes. Adding a fixture is adding a file; see `docs/fixtures.md`.
3
4use serde_json::{Map, Value};
5use yaml_rust2::{Yaml, YamlLoader};
6
7#[derive(Debug, Clone, PartialEq)]
8pub struct Fixture {
9    pub id: String,
10    pub description: String,
11    pub kind: Kind,
12    /// How many times to run; a fixture passes only if every run passes.
13    pub repeat: u32,
14    /// Counts toward a `partial` verdict.
15    pub core: bool,
16    pub steps: Vec<Step>,
17    pub pass: Cond,
18    pub fail: Vec<FailRule>,
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum Kind {
23    /// Chat requests checked with conditions.
24    Chat,
25    /// Compare the runtime's chat template with the expected one.
26    Template,
27}
28
29#[derive(Debug, Clone, PartialEq)]
30pub struct Step {
31    pub stream: bool,
32    /// Merged over `{model, temperature: 0, seed: 42, max_tokens: 512}`. Objects of the
33    /// form `{"$from": "N:<jsonpath>"}` are replaced by that value from exchange N.
34    pub body: Value,
35}
36
37#[derive(Debug, Clone, PartialEq)]
38pub struct FailRule {
39    pub class: String,
40    pub when: Cond,
41}
42
43/// Where a condition looks: exchange `N` (default: the last one) and a JSONPath.
44#[derive(Debug, Clone, PartialEq)]
45pub struct Ref {
46    pub exchange: Option<usize>,
47    pub path: String,
48}
49
50#[derive(Debug, Clone, PartialEq)]
51pub enum Cond {
52    All(Vec<Cond>),
53    Any(Vec<Cond>),
54    Not(Box<Cond>),
55    Path(Ref, Test),
56    /// HTTP status: exact `code`, or the hundreds `class` (2 for 2xx).
57    Status {
58        exchange: Option<usize>,
59        code: Option<u16>,
60        class: Option<u16>,
61    },
62    /// The stream ended with `[DONE]`.
63    StreamDone {
64        exchange: Option<usize>,
65        done: bool,
66    },
67    /// Both references resolve to equal values (strings holding JSON are compared parsed).
68    Same(Ref, Ref),
69    /// Exchange N happened (earlier steps produced what later ones needed).
70    Ran(usize),
71    Always,
72}
73
74#[derive(Debug, Clone, PartialEq)]
75pub enum Test {
76    Exists(bool),
77    Count {
78        min: Option<u64>,
79        max: Option<u64>,
80    },
81    Equals(Value),
82    ContainsAny(Vec<String>),
83    /// Every match is a string holding valid JSON (or not).
84    Json(bool),
85    /// Every match (parsed from a string when `parse`) validates against the schema.
86    Schema {
87        schema: Value,
88        parse: bool,
89    },
90}
91
92#[derive(Debug, thiserror::Error)]
93#[error("fixture {file}: {message}")]
94pub struct Error {
95    pub file: String,
96    pub message: String,
97}
98
99pub fn yaml_to_json(y: &Yaml) -> Value {
100    match y {
101        Yaml::Real(r) => r.parse::<f64>().map_or(Value::Null, |f| {
102            serde_json::Number::from_f64(f).map_or(Value::Null, Value::Number)
103        }),
104        Yaml::Integer(i) => Value::from(*i),
105        Yaml::String(s) => Value::String(s.clone()),
106        Yaml::Boolean(b) => Value::Bool(*b),
107        Yaml::Array(a) => Value::Array(a.iter().map(yaml_to_json).collect()),
108        Yaml::Hash(h) => {
109            let mut m = Map::new();
110            for (k, v) in h {
111                let key = match k {
112                    Yaml::String(s) => s.clone(),
113                    Yaml::Integer(i) => i.to_string(),
114                    Yaml::Boolean(b) => b.to_string(),
115                    _ => continue,
116                };
117                m.insert(key, yaml_to_json(v));
118            }
119            Value::Object(m)
120        }
121        _ => Value::Null,
122    }
123}
124
125pub fn parse(file: &str, text: &str) -> Result<Fixture, Error> {
126    let err = |message: String| Error {
127        file: file.to_owned(),
128        message,
129    };
130    let docs = YamlLoader::load_from_str(text).map_err(|e| err(format!("invalid YAML: {e}")))?;
131    let v = docs.first().map(yaml_to_json).unwrap_or(Value::Null);
132    let o = v
133        .as_object()
134        .ok_or_else(|| err("top level must be a mapping".into()))?;
135    let known = [
136        "id",
137        "description",
138        "kind",
139        "repeat",
140        "core",
141        "steps",
142        "pass",
143        "fail",
144    ];
145    if let Some(k) = o.keys().find(|k| !known.contains(&k.as_str())) {
146        return Err(err(format!("unknown key `{k}`")));
147    }
148    let id = o
149        .get("id")
150        .and_then(Value::as_str)
151        .ok_or_else(|| err("missing `id`".into()))?
152        .to_owned();
153    let kind = match o.get("kind").and_then(Value::as_str).unwrap_or("chat") {
154        "chat" => Kind::Chat,
155        "template" => Kind::Template,
156        k => return Err(err(format!("unknown kind `{k}`"))),
157    };
158    let steps = o
159        .get("steps")
160        .and_then(Value::as_array)
161        .map(|s| {
162            s.iter()
163                .map(|st| {
164                    Ok(Step {
165                        stream: st["stream"].as_bool().unwrap_or(false),
166                        body: st
167                            .get("body")
168                            .cloned()
169                            .ok_or_else(|| err("a step needs a `body`".into()))?,
170                    })
171                })
172                .collect::<Result<Vec<_>, Error>>()
173        })
174        .transpose()?
175        .unwrap_or_default();
176    if kind == Kind::Chat && steps.is_empty() {
177        return Err(err("a chat fixture needs `steps`".into()));
178    }
179    let pass = match o.get("pass") {
180        Some(c) => cond(c).map_err(|m| err(format!("pass: {m}")))?,
181        None => Cond::Always,
182    };
183    let fail = o
184        .get("fail")
185        .and_then(Value::as_array)
186        .map(|rules| {
187            rules
188                .iter()
189                .map(|r| {
190                    let class = r["class"]
191                        .as_str()
192                        .ok_or_else(|| err("a fail rule needs `class`".into()))?
193                        .to_owned();
194                    let when = match r.get("if") {
195                        Some(c) => cond(c).map_err(|m| err(format!("fail {class}: {m}")))?,
196                        None => Cond::Always,
197                    };
198                    Ok(FailRule { class, when })
199                })
200                .collect::<Result<Vec<_>, Error>>()
201        })
202        .transpose()?
203        .unwrap_or_default();
204    Ok(Fixture {
205        id,
206        description: o
207            .get("description")
208            .and_then(Value::as_str)
209            .unwrap_or("")
210            .trim()
211            .to_owned(),
212        kind,
213        repeat: o.get("repeat").and_then(Value::as_u64).unwrap_or(3) as u32,
214        core: o.get("core").and_then(Value::as_bool).unwrap_or(false),
215        steps,
216        pass,
217        fail,
218    })
219}
220
221/// `"1:$.a.b"` or `"$.a.b"`.
222pub fn parse_ref(s: &str) -> Result<Ref, String> {
223    let (exchange, path) = match s.split_once(':') {
224        Some((n, p)) if n.bytes().all(|b| b.is_ascii_digit()) && !n.is_empty() => (
225            Some(n.parse().map_err(|_| format!("bad exchange in `{s}`"))?),
226            p,
227        ),
228        _ => (None, s),
229    };
230    serde_json_path::JsonPath::parse(path).map_err(|e| format!("bad JSONPath `{path}`: {e}"))?;
231    Ok(Ref {
232        exchange,
233        path: path.to_owned(),
234    })
235}
236
237fn exchange_of(o: &Map<String, Value>) -> Option<usize> {
238    o.get("exchange")
239        .and_then(Value::as_u64)
240        .map(|n| n as usize)
241}
242
243fn cond(v: &Value) -> Result<Cond, String> {
244    let o = v.as_object().ok_or("a condition must be a mapping")?;
245    let list = |k: &str| -> Result<Vec<Cond>, String> {
246        o[k].as_array()
247            .ok_or(format!("`{k}` takes a list"))?
248            .iter()
249            .map(cond)
250            .collect()
251    };
252    if o.contains_key("all") {
253        return Ok(Cond::All(list("all")?));
254    }
255    if o.contains_key("any") {
256        return Ok(Cond::Any(list("any")?));
257    }
258    if let Some(n) = o.get("not") {
259        return Ok(Cond::Not(Box::new(cond(n)?)));
260    }
261    if let Some(n) = o.get("ran").and_then(Value::as_u64) {
262        return Ok(Cond::Ran(n as usize));
263    }
264    if o.contains_key("status") || o.contains_key("status_class") {
265        return Ok(Cond::Status {
266            exchange: exchange_of(o),
267            code: o.get("status").and_then(Value::as_u64).map(|c| c as u16),
268            class: o
269                .get("status_class")
270                .and_then(Value::as_u64)
271                .map(|c| c as u16),
272        });
273    }
274    if let Some(d) = o.get("stream_done").and_then(Value::as_bool) {
275        return Ok(Cond::StreamDone {
276            exchange: exchange_of(o),
277            done: d,
278        });
279    }
280    if let Some(pair) = o.get("same").and_then(Value::as_array) {
281        let [a, b] = pair.as_slice() else {
282            return Err("`same` takes two references".into());
283        };
284        return Ok(Cond::Same(
285            parse_ref(a.as_str().ok_or("`same` references are strings")?)?,
286            parse_ref(b.as_str().ok_or("`same` references are strings")?)?,
287        ));
288    }
289    let path = o.get("path").and_then(Value::as_str).ok_or(
290        "unknown condition (expected path, status, stream_done, same, ran, all, any or not)",
291    )?;
292    let r = parse_ref(path)?;
293    let test = if let Some(b) = o.get("exists").and_then(Value::as_bool) {
294        Test::Exists(b)
295    } else if let Some(n) = o.get("count").and_then(Value::as_u64) {
296        Test::Count {
297            min: Some(n),
298            max: Some(n),
299        }
300    } else if o.contains_key("count_min") || o.contains_key("count_max") {
301        Test::Count {
302            min: o.get("count_min").and_then(Value::as_u64),
303            max: o.get("count_max").and_then(Value::as_u64),
304        }
305    } else if let Some(e) = o.get("equals") {
306        Test::Equals(e.clone())
307    } else if let Some(c) = o.get("contains_any").and_then(Value::as_array) {
308        Test::ContainsAny(
309            c.iter()
310                .filter_map(|s| s.as_str().map(str::to_owned))
311                .collect(),
312        )
313    } else if let Some(j) = o.get("json").and_then(Value::as_bool) {
314        Test::Json(j)
315    } else if let Some(s) = o.get("schema") {
316        jsonschema::validator_for(s).map_err(|e| format!("invalid schema: {e}"))?;
317        Test::Schema {
318            schema: s.clone(),
319            parse: o.get("parse").and_then(Value::as_bool).unwrap_or(false),
320        }
321    } else {
322        return Err(format!(
323            "`{path}`: no test (exists, count, count_min, count_max, equals, contains_any, json, schema)"
324        ));
325    };
326    Ok(Cond::Path(r, test))
327}
328
329/// The built-in fixtures, embedded at build time from `fixtures/` (every `*.yaml`).
330pub fn builtin() -> Vec<Fixture> {
331    include!(concat!(env!("OUT_DIR"), "/fixtures.rs"))
332        .iter()
333        .map(|(name, text)| parse(name, text).unwrap_or_else(|e| panic!("built-in {e}")))
334        .collect()
335}
336
337/// Fixtures from a directory (every `*.yaml`, sorted by name).
338pub fn load_dir(dir: &std::path::Path) -> Result<Vec<Fixture>, Error> {
339    let mut files: Vec<_> = std::fs::read_dir(dir)
340        .map_err(|e| Error {
341            file: dir.display().to_string(),
342            message: e.to_string(),
343        })?
344        .filter_map(|e| e.ok().map(|e| e.path()))
345        .filter(|p| p.extension().is_some_and(|x| x == "yaml" || x == "yml"))
346        .collect();
347    files.sort();
348    files
349        .iter()
350        .map(|p| {
351            let text = std::fs::read_to_string(p).map_err(|e| Error {
352                file: p.display().to_string(),
353                message: e.to_string(),
354            })?;
355            parse(&p.display().to_string(), &text)
356        })
357        .collect()
358}
loadmaster_probe/fixture.rs

loadmaster_probe/
fixture.rs
