tests.rs

crates/cortex-core/src/storage/tests.rs

1. use super::*;
2. use crate::types::{Prices, Usage, UsageSource};
4. fn db() -> Db {
5. Db::open_in_memory().unwrap()
6. }
8. fn chat(db: &Db) -> Conversation {
9. db.create_conversation(DEEPINFRA_ID, "standin/echo", "be brief", None, "main").unwrap()
10. }
12. fn reply(db: &Db, conv: &str, text: &str, cost: i64, source: UsageSource) -> Message {
13. let m = db.insert_assistant_placeholder(conv, DEEPINFRA_ID, "standin/echo").unwrap();
14. db.finalize_assistant(
15. &m.id,
16. &Finalize {
17. content: text.into(),
18. status: "complete",
19. usage: Usage { prompt_tokens: 10, completion_tokens: 20, cached_tokens: 0 },
20. usage_source: Some(source),
21. prices: Some(Prices { input_per_m: 0.5, output_per_m: 1.5, cache_read_per_m: None }),
22. cost_nanousd: cost,
23. ..Default::default()
24. },
25. )
26. .unwrap()
27. }
29. #[test]
30. fn migrations_run_once_and_fts5_is_compiled_in() {
31. let d = db();
32. assert_eq!(d.schema_version().unwrap(), MIGRATIONS.len() as i64);
33. let fts: bool = d.conn().query_row("SELECT sqlite_compileoption_used('ENABLE_FTS5')", [], |r| r.get(0)).unwrap();
34. assert!(fts);
35. // DeepInfra is seeded exactly once.
36. let dir = tempfile::tempdir().unwrap();
37. let path = dir.path().join("c.db");
38. Db::open(&path).unwrap();
39. let again = Db::open(&path).unwrap();
40. assert_eq!(again.list_providers().unwrap().len(), 1);
41. assert_eq!(again.schema_version().unwrap(), MIGRATIONS.len() as i64);
42. }
44. #[test]
45. fn a_conversation_round_trip() {
46. let d = db();
47. let c = chat(&d);
48. assert_eq!(c.title, "New chat");
49. d.insert_user_message(&c.id, "How does FTS5 rank results?\nsecond line", &[]).unwrap();
50. let a = reply(&d, &c.id, "It uses bm25.", 35_000, UsageSource::Reported);
51. assert_eq!(a.status, "complete");
52. assert_eq!(a.prices.unwrap().output_per_m, 1.5);
53. let msgs = d.list_messages(&c.id).unwrap();
54. assert_eq!(msgs.iter().map(|m| m.role.as_str()).collect::<Vec<_>>(), ["user", "assistant"]);
55. // Titled from the first line of the first message, and only once.
56. assert_eq!(d.get_conversation(&c.id).unwrap().title, "How does FTS5 rank results?");
57. d.insert_user_message(&c.id, "something else", &[]).unwrap();
58. assert_eq!(d.get_conversation(&c.id).unwrap().title, "How does FTS5 rank results?");
59. let list = d.list_conversations(false, 50).unwrap();
60. assert_eq!(list[0].cost_nanousd, 35_000);
61. assert_eq!(list[0].message_count, 3);
62. assert!(list[0].has_usage);
63. }
65. #[test]
66. fn a_chat_whose_only_reply_failed_is_not_free() {
67. let d = db();
68. let c = chat(&d);
69. d.insert_user_message(&c.id, "hi", &[]).unwrap();
70. let p = d.insert_assistant_placeholder(&c.id, DEEPINFRA_ID, "m").unwrap();
71. d.finalize_assistant(
72. &p.id,
73. &Finalize { status: "error", error_kind: Some("invalid_key".into()), ..Default::default() },
74. )
75. .unwrap();
76. let s = &d.list_conversations(false, 10).unwrap()[0];
77. assert_eq!((s.cost_nanousd, s.has_usage), (0, false));
78. }
80. #[test]
81. fn title_trimming() {
82. assert_eq!(title_from("\n\n hello \nworld"), "hello");
83. let long = "word ".repeat(30);
84. let t = title_from(&long);
85. assert!(t.ends_with("...") && t.chars().count() <= 63, "{t}");
86. assert_eq!(title_from(""), "New chat");
87. }
89. #[test]
90. fn truncate_from_removes_that_message_and_later() {
91. let d = db();
92. let c = chat(&d);
93. let u1 = d.insert_user_message(&c.id, "one", &[]).unwrap();
94. reply(&d, &c.id, "r1", 1, UsageSource::Reported);
95. let u2 = d.insert_user_message(&c.id, "two", &[]).unwrap();
96. reply(&d, &c.id, "r2", 1, UsageSource::Reported);
97. assert_eq!(d.truncate_from(&c.id, &u2.id).unwrap(), 2);
98. let left: Vec<String> = d.list_messages(&c.id).unwrap().into_iter().map(|m| m.content).collect();
99. assert_eq!(left, ["one", "r1"]);
100. assert_eq!(d.truncate_from(&c.id, &u1.id).unwrap(), 2);
101. }
103. #[test]
104. fn editing_a_message_keeps_its_attachments_for_the_resend() {
105. let d = db();
106. let c = chat(&d);
107. let a = d
108. .insert_attachment(&NewAttachment {
109. kind: "text",
110. filename: "a.txt",
111. mime: "text/plain",
112. byte_size: 1,
113. sha256: "x",
114. width: None,
115. height: None,
116. path: None,
117. text_content: Some("a"),
118. })
119. .unwrap();
120. let u = d.insert_user_message(&c.id, "see file", std::slice::from_ref(&a.id)).unwrap();
121. d.truncate_from(&c.id, &u.id).unwrap();
122. let again = d.insert_user_message(&c.id, "see this file", std::slice::from_ref(&a.id)).unwrap();
123. assert_eq!(again.attachments.len(), 1, "the attachment survives the edit");
124. }
126. #[test]
127. fn deleting_a_conversation_cascades() {
128. let d = db();
129. let c = chat(&d);
130. let a = d
131. .insert_attachment(&NewAttachment {
132. kind: "text",
133. filename: "notes.txt",
134. mime: "text/plain",
135. byte_size: 5,
136. sha256: "x",
137. width: None,
138. height: None,
139. path: None,
140. text_content: Some("hello"),
141. })
142. .unwrap();
143. let u = d.insert_user_message(&c.id, "see attached", std::slice::from_ref(&a.id)).unwrap();
144. assert_eq!(u.attachments.len(), 1);
145. d.delete_conversation(&c.id).unwrap();
146. let n: i64 = d.conn().query_row("SELECT COUNT(*) FROM messages", [], |r| r.get(0)).unwrap();
147. let na: i64 = d.conn().query_row("SELECT COUNT(*) FROM attachments", [], |r| r.get(0)).unwrap();
148. assert_eq!((n, na), (0, 0));
149. // The index follows: nothing left to find.
150. assert!(d.search("attached", 10).unwrap().is_empty());
151. }
153. #[test]
154. fn interrupted_streams_are_closed_on_open() {
155. let dir = tempfile::tempdir().unwrap();
156. let path = dir.path().join("c.db");
157. let d = Db::open(&path).unwrap();
158. let c = chat(&d);
159. let p = d.insert_assistant_placeholder(&c.id, DEEPINFRA_ID, "m").unwrap();
160. drop(d);
161. let d = Db::open(&path).unwrap();
162. let m = d.get_message(&p.id).unwrap();
163. assert_eq!((m.status.as_str(), m.error_kind.as_deref()), ("error", Some("interrupted")));
164. }
166. #[test]
167. fn search_finds_messages_and_titles_and_survives_odd_input() {
168. let d = db();
169. let c = chat(&d);
170. d.insert_user_message(&c.id, "Tidy a bash backup script", &[]).unwrap();
171. reply(&d, &c.id, "Use rsync with --delete and quote every variable.", 1, UsageSource::Reported);
172. let other = chat(&d);
173. d.insert_user_message(&other.id, "Weekend sourdough schedule", &[]).unwrap();
175. let hits = d.search("rsync", 10).unwrap();
176. assert_eq!(hits.len(), 1);
177. assert_eq!(hits[0].conversation_id, c.id);
178. assert!(hits[0].snippet.contains("\u{1}rsync\u{2}"), "{}", hits[0].snippet);
179. // Prefix match on the last word while typing; porter stemming on the rest.
180. assert!(!d.search("sourd", 10).unwrap().is_empty());
181. assert!(!d.search("quoting variables", 10).unwrap().is_empty());
182. // Operators and quotes are just text.
183. for q in ["\"", "AND OR NOT", "a:b", "(", "*", "-", "' OR 1=1 --", ""] {
184. d.search(q, 10).unwrap();
185. }
186. // Edited titles are found.
187. d.rename_conversation(&other.id, "Bread plans").unwrap();
188. assert!(d.search("bread", 10).unwrap().iter().any(|h| h.conversation_id == other.id && h.message_id.is_none()));
189. // Archived chats drop out of search.
190. d.set_archived(&c.id, true).unwrap();
191. assert!(d.search("rsync", 10).unwrap().is_empty());
192. }
194. #[test]
195. fn fts_query_sanitizer() {
196. assert_eq!(fts_query("hello world").unwrap(), "\"hello\" \"world\"*");
197. assert_eq!(fts_query(" (a:b) \"c ").unwrap(), "\"a\" \"b\" \"c\"*");
198. assert_eq!(fts_query("don't"), Some("\"don't\"*".into()));
199. assert_eq!(fts_query("!!! ---"), None);
200. }
202. #[test]
203. fn spend_and_month_boundaries() {
204. let d = db();
205. let c = chat(&d);
206. let old = reply(&d, &c.id, "old", 5_000, UsageSource::Reported);
207. reply(&d, &c.id, "new", 7_000, UsageSource::Estimated);
208. // Move the first reply to 2026-08-31 23:30 in UTC-5, which is September in UTC.
209. let utc_minus_5 = chrono::FixedOffset::west_opt(5 * 3600).unwrap();
210. let aug_local = utc_minus_5.with_ymd_and_hms(2026, 8, 31, 23, 30, 0).unwrap();
211. d.conn()
212. .execute("UPDATE messages SET created_at = ?1 WHERE id = ?2", params![aug_local.timestamp_millis(), old.id])
213. .unwrap();
214. let sept_local = utc_minus_5.with_ymd_and_hms(2026, 9, 15, 12, 0, 0).unwrap();
215. let start = month_start_ms(&sept_local);
216. assert_eq!(start, utc_minus_5.with_ymd_and_hms(2026, 9, 1, 0, 0, 0).unwrap().timestamp_millis());
217. let (spend, estimated) = d.spend_since(start).unwrap();
218. assert_eq!((spend, estimated), (7_000, true), "an August-local reply must not count in September");
219. // In UTC the same instant is already September.
220. let sept_utc = chrono::Utc.with_ymd_and_hms(2026, 9, 15, 12, 0, 0).unwrap();
221. assert_eq!(d.spend_since(month_start_ms(&sept_utc)).unwrap().0, 12_000);
222. // Grouped views.
223. let by_model = d.usage_by_model(0).unwrap();
224. assert_eq!((by_model[0].key.as_str(), by_model[0].cost_nanousd, by_model[0].messages), ("standin/echo", 12_000, 2));
225. let by_day = d.usage_by_day(0, -300).unwrap();
226. assert!(by_day.iter().any(|r| r.key == "2026-08-31" && r.cost_nanousd == 5_000));
227. }
229. #[test]
230. fn settings_presets_models() {
231. let d = db();
232. assert_eq!(d.get_setting::<f64>("budget").unwrap(), None);
233. d.set_setting("budget", &12.5).unwrap();
234. assert_eq!(d.get_setting::<f64>("budget").unwrap(), Some(12.5));
236. let p = d
237. .save_preset(&Preset {
238. id: String::new(),
239. name: "Editor".into(),
240. system_prompt: "Tighten prose.".into(),
241. provider_id: None,
242. model_id: None,
243. temperature: Some(0.3),
244. max_tokens: None,
245. })
246. .unwrap();
247. d.save_preset(&Preset { system_prompt: "Tighten prose. Prefer short sentences.".into(), ..p.clone() }).unwrap();
248. let all = d.list_presets().unwrap();
249. assert_eq!(all.len(), 1);
250. assert!(all[0].system_prompt.ends_with("Prefer short sentences."));
252. let models = crate::pricing::parse_priced_models(DEEPINFRA_ID, &standin::models());
253. d.replace_models(DEEPINFRA_ID, &models, "live", 1000).unwrap();
254. assert_eq!(d.list_models(DEEPINFRA_ID).unwrap(), models);
255. assert_eq!(d.models_source(DEEPINFRA_ID).unwrap(), Some(("live".into(), 1000)));
256. d.replace_models(DEEPINFRA_ID, &models[..2], "snapshot", 2000).unwrap();
257. assert_eq!(d.list_models(DEEPINFRA_ID).unwrap().len(), 2);
258. }
260. #[test]
261. fn providers() {
262. let d = db();
263. d.upsert_provider(&Provider {
264. id: "ollama".into(),
265. kind: "openai_compatible".into(),
266. name: "Ollama".into(),
267. base_url: "http://localhost:11434/v1/".into(),
268. requires_key: false,
269. })
270. .unwrap();
271. let all = d.list_providers().unwrap();
272. assert_eq!(all[0].id, DEEPINFRA_ID);
273. assert_eq!(all[1].base_url, "http://localhost:11434/v1");
274. d.delete_provider(DEEPINFRA_ID).unwrap();
275. d.delete_provider("ollama").unwrap();
276. assert_eq!(d.list_providers().unwrap().len(), 1);
277. }