1use std::path::Path;
2
3use sqlite::Connection;
4
5#[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
19pub 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
83fn 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}