mod.rsannotatedmod.rssource85 lines · 3.0 KB · raw
1pub mod dns;
2pub mod gates;
3pub mod postgres;
4
5use std::path::Path;
6
7pub use jev_mock as mock_jev;
8
9/// A Postgres serving the built extension, pointed at `mock`, with one
10/// row in `tickets` and a statement timeout so a hang fails fast.
11pub async fn jev_instance(mock: &mock_jev::MockJev) -> (postgres::Instance, tokio_postgres::Client) {
12    jev_instance_with(mock, &[]).await
13}
14
15/// As [`jev_instance`], with extra server settings (later ones win).
16pub async fn jev_instance_with(
17    mock: &mock_jev::MockJev,
18    settings: &[(&str, &str)],
19) -> (postgres::Instance, tokio_postgres::Client) {
20    jev_instance_env(mock, settings, &[("TYPESAFE_API_KEY", "test-key")]).await
21}
22
23/// As [`jev_instance_with`], with the server's environment given in full
24/// (no `TYPESAFE_API_KEY` unless `env` sets it).
25pub async fn jev_instance_env(
26    mock: &mock_jev::MockJev,
27    settings: &[(&str, &str)],
28    env: &[(&str, &str)],
29) -> (postgres::Instance, tokio_postgres::Client) {
30    let mut pg = postgres::TestPostgres::new()
31        .with_extension(Path::new(&std::env::var("POSTJEVSQL_EXT").unwrap()))
32        .setting("jev.endpoint", &mock.endpoint())
33        .setting("jev.model", "jev-1.13.0")
34        .setting("jev.ca_file", mock.ca_file().to_str().unwrap());
35    for (name, value) in env {
36        pg = pg.env_var(name, value);
37    }
38    for (name, value) in settings {
39        pg = pg.setting(name, value);
40    }
41    let pg = pg.start().await;
42    let (client, conn) = tokio_postgres::connect(&pg.conn_str(), tokio_postgres::NoTls)
43        .await
44        .expect("connects");
45    tokio::spawn(conn);
46    client
47        .batch_execute(&format!(
48            "SET statement_timeout = '20s';
49             CREATE EXTENSION postjevsql;
50             CREATE TABLE tickets (id int, body text);
51             INSERT INTO tickets VALUES (1, '{TICKET}');"
52        ))
53        .await
54        .expect("setup");
55    (pg, client)
56}
57
58pub const TICKET: &str = "Help! My payouts have been failing for 3 days.";
59
60pub fn noul(p: f64) -> serde_json::Value {
61    serde_json::json!({
62        "model": "jev-1.13.0",
63        "answers": { "q": { "type": "noul", "noul": p } },
64        "usage": { "input_tokens": 296, "output_tokens": 20 }
65    })
66}
67
68/// Compares two JSON values by what their numbers are worth. Postgres
69/// renders a float8 1.0 as `1` in `to_jsonb`, while a jsonb built in Rust
70/// keeps `1.0`, and serde_json holds those unequal; a record's fields
71/// are compared as values, not as how they were spelled.
72#[track_caller]
73pub fn assert_json_eq(actual: &serde_json::Value, expected: &serde_json::Value) {
74    assert_eq!(by_value(actual), by_value(expected));
75}
76
77fn by_value(v: &serde_json::Value) -> serde_json::Value {
78    use serde_json::Value;
79    match v {
80        Value::Number(n) => n.as_f64().and_then(serde_json::Number::from_f64).map_or_else(|| v.clone(), Value::Number),
81        Value::Array(a) => Value::Array(a.iter().map(by_value).collect()),
82        Value::Object(o) => Value::Object(o.iter().map(|(k, v)| (k.clone(), by_value(v))).collect()),
83        _ => v.clone(),
84    }
85}