5use std::sync::atomic::{AtomicUsize, Ordering};
6use std::sync::{Arc, Mutex};

8use serde_json::json;
9use support::mock_jev::{MockJev, Reply};
10use support::{jev_instance, noul};
11use tokio_postgres::AsyncMessage;
12
13const STATS: &str = "SELECT requests, retries, redials, connections, cache_hits, cache_misses,
14                            dedupe_hits, input_tokens, output_tokens, cost, in_flight FROM jev_stats()";
15
16#[derive(Debug, PartialEq)]
17struct Stats {
18    requests: i64,
19    retries: i64,
20    redials: i64,
21    connections: i64,
22    cache_hits: i64,
23    cache_misses: i64,
24    dedupe_hits: i64,
25    input_tokens: i64,
26    output_tokens: i64,
27    cost: f64,
28    in_flight: i64,
29}
30
31async fn stats(client: &tokio_postgres::Client) -> Stats {
32    let r = client.query_one(STATS, &[]).await.expect("jev_stats");
33    Stats {
34        requests: r.get(0),
35        retries: r.get(1),
36        redials: r.get(2),
37        connections: r.get(3),
38        cache_hits: r.get(4),
39        cache_misses: r.get(5),
40        dedupe_hits: r.get(6),
41        input_tokens: r.get(7),
42        output_tokens: r.get(8),
43        cost: r.get(9),
44        in_flight: r.get(10),
45    }
46}
47
48fn zero() -> Stats {
49    Stats {
50        requests: 0,
51        retries: 0,
52        redials: 0,
53        connections: 0,
54        cache_hits: 0,
55        cache_misses: 0,
56        dedupe_hits: 0,
57        input_tokens: 0,
58        output_tokens: 0,
59        cost: 0.0,
60        in_flight: 0,
61    }
62}
63
64#[tokio::test(flavor = "multi_thread")]
65async fn a_fresh_session_has_done_nothing() {
66    let mock = MockJev::start(|_| Reply::json(200, noul(0.95))).await;
67    let (_pg, client) = jev_instance(&mock).await;
68    assert_eq!(stats(&client).await, zero());
69}
70
71#[tokio::test(flavor = "multi_thread")]
72async fn attempts_retries_tokens_and_spend_are_counted() {
73    let calls = AtomicUsize::new(0);
74    let mock = MockJev::start(move |_| match calls.fetch_add(1, Ordering::SeqCst) {
75        0 => Reply::json(503, json!({ "detail": { "error_type": "unavailable", "message": "down" } }))
76            .header("retry-after-ms", "10"),
77        _ => Reply::json(200, noul(0.95)),
78    })
79    .await;
80    let (_pg, client) = jev_instance(&mock).await;
81    client.batch_execute("SET jev.price_per_mtok = 1000").await.unwrap();
82
83    client.query_one("SELECT jev_prob(t, 'Urgent?') FROM tickets t", &[]).await.expect("retried to success");
84
85    let s = stats(&client).await;
86    assert_eq!(
87        s,
88        Stats {
89            requests: 2,
90            retries: 1,
91            connections: 1,
92            cache_misses: 1,
93            input_tokens: 296,
94            output_tokens: 20,
95            // Input only; output tokens are not billed.
96            cost: 296.0 * 1000.0 / 1e6,
97            ..zero()
98        }
99    );
100}
101
102#[tokio::test(flavor = "multi_thread")]
103async fn a_redial_is_counted_apart_from_retries() {
104    let mock = MockJev::start(|_| Reply::json(200, noul(0.95))).await;
105    let (_pg, client) = jev_instance(&mock).await;
106    client.query_one("SELECT jev_prob(t, 'Urgent?') FROM tickets t", &[]).await.expect("first");
107    mock.goaway();
108    tokio::time::sleep(std::time::Duration::from_millis(200)).await;
109    // A new question, so the cache cannot answer it.
110    client.query_one("SELECT jev_prob(t, 'Spam?') FROM tickets t", &[]).await.expect("second");
111
112    let s = stats(&client).await;
113    assert_eq!((s.requests, s.retries, s.redials, s.connections), (2, 0, 1, 2), "{s:?}");
114}
115
116#[tokio::test(flavor = "multi_thread")]
117async fn retries_and_answers_are_logged_at_debug_with_their_request_ids() {
118    let calls = AtomicUsize::new(0);
119    let mock = MockJev::start(move |_| match calls.fetch_add(1, Ordering::SeqCst) {
120        0 => Reply::json(429, json!({ "detail": { "error_type": "rate_limited", "message": "slow" } }))
121            .header("retry-after-ms", "10")
122            .header("x-typesafe-request-id", "req-429"),
123        _ => Reply::json(200, noul(0.95)).header("x-typesafe-request-id", "req-ok"),
124    })
125    .await;
126    let (pg, _setup) = jev_instance(&mock).await;
127
128    let (client, mut conn) = tokio_postgres::connect(&pg.conn_str(), tokio_postgres::NoTls).await.unwrap();
129    let notices = Arc::new(Mutex::new(Vec::<String>::new()));
130    let sink = notices.clone();
131    tokio::spawn(async move {
132        while let Some(m) = std::future::poll_fn(|cx| conn.poll_message(cx)).await {
133            if let Ok(AsyncMessage::Notice(n)) = m {
134                sink.lock().unwrap().push(n.message().to_owned());
135            }
136        }
137    });
138    client.batch_execute("SET statement_timeout = '20s'; SET client_min_messages = debug1").await.unwrap();
139    client.query_one("SELECT jev_prob(t, 'Urgent?') FROM tickets t", &[]).await.expect("retried");
140
141    let notices = notices.lock().unwrap().clone();
142    let jev: Vec<_> = notices.iter().filter(|n| n.starts_with("jev:")).collect();
143    assert!(
144        jev.iter().any(|n| n.contains("connected to")),
145        "{notices:#?}"
146    );
147    assert!(
148        jev.iter().any(|n| n.contains("retrying") && n.contains("429") && n.contains("req-429")),
149        "{notices:#?}"
150    );
151    assert!(jev.iter().any(|n| n.contains("answered") && n.contains("req-ok")), "{notices:#?}");
152}
153
154#[tokio::test(flavor = "multi_thread")]
155async fn a_failed_request_leaves_the_in_flight_count() {
156    let mock = MockJev::start(|_| {
157        Reply::json(403, json!({ "detail": { "error_type": "forbidden", "message": "no key" } }))
158    })
159    .await;
160    let (_pg, client) = jev_instance(&mock).await;
161    client.query_one("SELECT jev_prob(t, 'Urgent?') FROM tickets t", &[]).await.expect_err("403");
162    let s = stats(&client).await;
163    assert_eq!((s.requests, s.in_flight), (1, 0), "{s:?}");
164}