7use serde_json::{Value, json};
8use tokio_postgres::error::SqlState;
9use support::jev_instance;
10use support::mock_jev::{MockJev, Reply};

15async fn first_label(fail_2: bool) -> MockJev {
16    MockJev::start(move |req| {
17        let body: Value = serde_json::from_slice(&req.body).unwrap();
18        if fail_2 && body["state"]["id"] == 2 {
19            return Reply::json(500, json!({ "detail": { "error_type": "server_error", "message": "boom" } }));
20        }
21        let mut answers = serde_json::Map::new();
22        for (qid, q) in body["questions"].as_object().unwrap() {
23            let answer = if q["type"] == "score" {
24                let n = q["criteria"].as_array().unwrap().len();
25                let ps: serde_json::Map<_, _> =
26                    (0..n).map(|i| (i.to_string(), json!(if i == 0 { 0.6 } else { 0.4 / (n - 1) as f64 }))).collect();
27                let legend: serde_json::Map<_, _> = (0..n).map(|i| (i.to_string(), q["criteria"][i].clone())).collect();
28                json!({ "type": "score", "score": 0.4, "confidence": 0.37, "legend": legend, "probabilities": ps })
29            } else {
30                let labels: Vec<&String> = q["criteria"].as_object().unwrap().keys().collect();
31                let ps: serde_json::Map<_, _> = labels
32                    .iter()
33                    .enumerate()
34                    .map(|(i, l)| (l.to_string(), json!(if i == 0 { 0.6 } else { 0.4 / (labels.len() - 1) as f64 })))
35                    .collect();
36                json!({ "type": "choice", "choice": labels[0], "confidence": 0.37, "probabilities": ps })
37            };
38            answers.insert(qid.clone(), answer);
39        }
40        Reply::json(200, json!({ "model": "jev-1.13.0", "answers": answers, "usage": { "input_tokens": 296, "output_tokens": 20 } }))
41    })
42    .await
43}

45const OPTIONS: &str = "ARRAY['sales', 'billing', 'other']";
46
47#[tokio::test(flavor = "multi_thread")]
48async fn shares_the_choice_judgment_and_its_cache_row() {
49    let mock = first_label(false).await;
50    let (_pg, client) = jev_instance(&mock).await;
51
52    let sql = format!(
53        "SELECT jev_choice(t, 'Which team?', {OPTIONS}), jev_confidence(t, 'Which team?', 'choice', {OPTIONS}) FROM tickets t"
54    );
55    let row = client.query_one(&sql, &[]).await.unwrap();
56    assert_eq!((row.get::<_, String>(0), row.get::<_, f64>(1)), ("sales".to_owned(), 0.37), "as the vendor returned it");
57    assert_eq!(mock.requests().len(), 1, "one request for both calls");
58    let sent: Value = serde_json::from_slice(&mock.requests()[0].body).unwrap();
59    assert_eq!(sent["questions"].as_object().unwrap().len(), 1, "one question for both calls");
60    let cached: i64 = client.query_one("SELECT count(*) FROM jev_cache", &[]).await.unwrap().get(0);
61    assert_eq!(cached, 1);
62
63    // A rerun of either is a cache hit, and sends nothing.
64    let before: i64 = client.query_one("SELECT cache_hits FROM jev_stats()", &[]).await.unwrap().get(0);
65    let confidence: f64 = client
66        .query_one(&format!("SELECT jev_confidence(t, 'Which team?', 'choice', {OPTIONS}) FROM tickets t"), &[])
67        .await
68        .unwrap()
69        .get(0);
70    assert_eq!(confidence, 0.37);
71    let after: i64 = client.query_one("SELECT cache_hits FROM jev_stats()", &[]).await.unwrap().get(0);
72    assert_eq!(after - before, 1);
73    assert_eq!(mock.requests().len(), 1, "served from the cache");
74}
75
76#[tokio::test(flavor = "multi_thread")]
77async fn shares_the_score_judgment() {
78    let mock = first_label(false).await;
79    let (_pg, client) = jev_instance(&mock).await;
80
81    let levels = "ARRAY['calm', 'annoyed', 'furious']";
82    let row = client
83        .query_one(
84            &format!("SELECT jev_score(t, 'How upset?', {levels}), jev_confidence(t, 'How upset?', 'score', {levels}) FROM tickets t"),
85            &[],
86        )
87        .await
88        .unwrap();
89    assert_eq!((row.get::<_, f64>(0), row.get::<_, f64>(1)), (0.4, 0.37));
90    assert_eq!(mock.requests().len(), 1);
91}
92
93#[tokio::test(flavor = "multi_thread")]
94async fn unsure_is_null() {
95    let mock = first_label(true).await;
96    let (_pg, client) = jev_instance(&mock).await;
97    client.batch_execute("INSERT INTO tickets VALUES (2, 'Where is my refund?'); SET jev.on_error = unsure").await.unwrap();
98
99    let rows = client
100        .query(&format!("SELECT id, jev_confidence(t, 'Which team?', 'choice', {OPTIONS}) FROM tickets t ORDER BY id"), &[])
101        .await
102        .expect("a failed judgment does not fail the statement");
103    let got: Vec<(i32, Option<f64>)> = rows.iter().map(|r| (r.get(0), r.get(1))).collect();
104    assert_eq!(got, [(1, Some(0.37)), (2, None)]);
105}
106
107#[tokio::test(flavor = "multi_thread")]
108async fn a_noul_has_no_confidence() {
109    let mock = first_label(false).await;
110    let (_pg, client) = jev_instance(&mock).await;
111
112    let err = client
113        .query_one("SELECT jev_confidence(t, 'Urgent?', 'noul', ARRAY[]::text[]) FROM tickets t", &[])
114        .await
115        .expect_err("refused, not derived from the probability");
116    let db = err.as_db_error().expect("an ERROR");
117    assert_eq!(db.code(), &SqlState::INVALID_PARAMETER_VALUE);
118    assert!(db.message().contains("carries no confidence"), "{}", db.message());
119    assert!(mock.requests().is_empty(), "nothing sent");
120}