Skip to main content

csusage_test_support/
zcode.rs

1use std::path::Path;
2
3use sqlite::Connection;
4
5/// Describes one completed ZCode usage row for a SQLite fixture.
6#[derive(Clone, Copy)]
7struct FixtureUsage<'a> {
8    id: &'a str,
9    session_id: &'a str,
10    timestamp: &'a str,
11    model: &'a str,
12    input_tokens: i64,
13    output_tokens: i64,
14    cache_creation_tokens: i64,
15    cache_read_tokens: i64,
16    computed_total_tokens: i64,
17}
18
19/// Creates the schema and representative rows used by ZCode report tests.
20pub fn create_fixture(path: impl AsRef<Path>) {
21    let db = sqlite::open(path).unwrap();
22    db.execute(
23        "CREATE TABLE model_usage (
24            id TEXT PRIMARY KEY, session_id TEXT, started_at INTEGER, model_id TEXT,
25            provider_id TEXT, status TEXT, input_tokens INTEGER, output_tokens INTEGER,
26            cache_creation_input_tokens INTEGER, cache_read_input_tokens INTEGER,
27            computed_total_tokens INTEGER
28        )",
29    )
30    .unwrap();
31    db.execute("CREATE TABLE session (id TEXT PRIMARY KEY, directory TEXT, version TEXT)")
32        .unwrap();
33    db.execute(
34        "INSERT INTO session VALUES
35            ('session-a', '/workspace/project-a', '0.16.3'),
36            ('session-b', '/workspace/project-b', '0.16.3')",
37    )
38    .unwrap();
39    insert_usage(
40        &db,
41        FixtureUsage {
42            id: "usage-52",
43            session_id: "session-a",
44            timestamp: "2099-01-02T00:00:00.000Z",
45            model: "GLM-5.2",
46            input_tokens: 100,
47            output_tokens: 10,
48            cache_creation_tokens: 15,
49            cache_read_tokens: 25,
50            computed_total_tokens: 110,
51        },
52    );
53    insert_usage(
54        &db,
55        FixtureUsage {
56            id: "usage-53",
57            session_id: "session-a",
58            timestamp: "2099-01-15T12:00:00.000Z",
59            model: "GLM-5.3",
60            input_tokens: 200,
61            output_tokens: 20,
62            cache_creation_tokens: 30,
63            cache_read_tokens: 40,
64            computed_total_tokens: 220,
65        },
66    );
67    insert_usage(
68        &db,
69        FixtureUsage {
70            id: "usage-53-b",
71            session_id: "session-b",
72            timestamp: "2099-02-01T00:00:00.000Z",
73            model: "GLM-5.3",
74            input_tokens: 50,
75            output_tokens: 5,
76            cache_creation_tokens: 0,
77            cache_read_tokens: 10,
78            computed_total_tokens: 55,
79        },
80    );
81}
82
83/// Inserts one completed usage row into a ZCode fixture database.
84fn insert_usage(db: &Connection, usage: FixtureUsage<'_>) {
85    let mut statement = db
86        .prepare(
87            "INSERT INTO model_usage
88             (id, session_id, started_at, model_id, provider_id, status, input_tokens,
89              output_tokens, cache_creation_input_tokens, cache_read_input_tokens,
90              computed_total_tokens)
91             VALUES (?1, ?2, ?3, ?4, 'builtin:zai-coding-plan', 'completed', ?5, ?6, ?7, ?8, ?9)",
92        )
93        .unwrap();
94    statement.bind((1, usage.id)).unwrap();
95    statement.bind((2, usage.session_id)).unwrap();
96    statement
97        .bind((
98            3,
99            usage
100                .timestamp
101                .parse::<jiff::Timestamp>()
102                .unwrap()
103                .as_millisecond(),
104        ))
105        .unwrap();
106    statement.bind((4, usage.model)).unwrap();
107    statement.bind((5, usage.input_tokens)).unwrap();
108    statement.bind((6, usage.output_tokens)).unwrap();
109    statement.bind((7, usage.cache_creation_tokens)).unwrap();
110    statement.bind((8, usage.cache_read_tokens)).unwrap();
111    statement.bind((9, usage.computed_total_tokens)).unwrap();
112    statement.next().unwrap();
113}