firefox-main/browser/components/aiwindow/ui/test/xpcshell/test_ConversationStore.js (file symbol)

Source code

Revision control

Copy as Markdown

Other Tools

Test Info: Warnings

/* Any copyright is dedicated to the Public Domain.
do_get_profile();
const { ConversationStore } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/ui/modules/ConversationStore.sys.mjs"
);
const { Conversation, MESSAGE_ROLE } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/models/Conversation.sys.mjs"
);
const { SecurityProperties } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/models/SecurityProperties.sys.mjs"
);
async function countRows(convId) {
const rows = await ConversationStore.connection.execute(
"SELECT count(*) AS c FROM conversation WHERE conv_id = :conv_id",
{ conv_id: convId }
);
return rows[0].getResultByName("c");
}
registerCleanupFunction(async () => {
await ConversationStore.destroyDatabase();
});
add_task(async function test_database_initializes() {
await ConversationStore.ensureDatabase();
Assert.ok(ConversationStore.connection, "A database connection is opened");
});
add_task(async function test_save_and_restore() {
const conversation = new Conversation({
id: "conv-save-restore",
createdDate: 1000,
updatedDate: 2000,
feature: "aitab",
serpUrlsForAnonymousFetch: ["https://serp.example/"],
// Non-default flags so the round-trip proves the values actually persist,
// rather than both sides defaulting to false.
securityProperties: { privateData: true, untrustedInput: true },
});
await ConversationStore.updateConversation(conversation);
const loaded =
await ConversationStore.findConversationById("conv-save-restore");
Assert.ok(loaded instanceof Conversation, "A Conversation is returned");
Assert.equal(loaded.id, "conv-save-restore", "id is restored");
Assert.equal(loaded.createdDate, 1000, "createdDate is restored");
Assert.equal(loaded.updatedDate, 2000, "updatedDate is restored");
Assert.equal(loaded.feature, "aitab", "feature is restored");
// A Set made inside a .sys.mjs module isn't `instanceof Set` from this test
// (different JS globals => different Set constructors). Check size/has instead.
Assert.equal(loaded.seenUrls.size, 2, "seenUrls restored with both values");
Assert.ok(
loaded.seenUrls.has("https://a.example/") &&
loaded.seenUrls.has("https://b.example/"),
"seenUrls values are restored"
);
Assert.equal(
loaded.serpUrlsForAnonymousFetch.size,
1,
"serpUrlsForAnonymousFetch restored"
);
Assert.ok(
loaded.serpUrlsForAnonymousFetch.has("https://serp.example/"),
"serpUrlsForAnonymousFetch values are restored"
);
Assert.ok(
loaded.securityProperties instanceof SecurityProperties,
"securityProperties is restored as a SecurityProperties"
);
Assert.ok(
loaded.securityProperties.privateData,
"privateData flag survives the round-trip"
);
Assert.ok(
loaded.securityProperties.untrustedInput,
"untrustedInput flag survives the round-trip"
);
});
add_task(async function test_update_is_upsert() {
await ConversationStore.updateConversation(
new Conversation({
id: "conv-upsert",
createdDate: 1000,
updatedDate: 2000,
feature: "chat",
})
);
// Same id, different created/updated/feature. created_date is insert-only.
await ConversationStore.updateConversation(
new Conversation({
id: "conv-upsert",
createdDate: 9999,
updatedDate: 3000,
feature: "aitab",
})
);
Assert.equal(await countRows("conv-upsert"), 1, "Upsert keeps a single row");
const loaded = await ConversationStore.findConversationById("conv-upsert");
Assert.equal(loaded.createdDate, 1000, "created_date is insert-only");
Assert.equal(loaded.updatedDate, 3000, "updated_date is refreshed");
Assert.equal(loaded.feature, "aitab", "feature is refreshed");
});
add_task(async function test_delete() {
await ConversationStore.updateConversation(
new Conversation({ id: "conv-delete", createdDate: 1, updatedDate: 1 })
);
Assert.equal(await countRows("conv-delete"), 1, "Row exists before delete");
await ConversationStore.deleteConversationById("conv-delete");
Assert.equal(await countRows("conv-delete"), 0, "Row is deleted");
Assert.equal(
await ConversationStore.findConversationById("conv-delete"),
null,
"findConversationById returns null after delete"
);
});
add_task(async function test_messages_save_and_restore() {
const conversation = new Conversation({
id: "conv-msgs",
createdDate: 100,
updatedDate: 200,
});
conversation.addMessage(MESSAGE_ROLE.USER, { type: "text", body: "hi" }, 0);
conversation.addMessage(
MESSAGE_ROLE.ASSISTANT,
{ type: "text", body: "hello" },
0
);
await ConversationStore.updateConversation(conversation);
const loaded = await ConversationStore.findConversationById("conv-msgs");
Assert.equal(loaded.messages.length, 2, "both messages are restored");
Assert.equal(
loaded.messages[0].role,
MESSAGE_ROLE.USER,
"the first message role is restored"
);
Assert.deepEqual(
loaded.messages[0].content,
{ type: "text", body: "hi" },
"message content is restored"
);
Assert.equal(
loaded.messages[1].content.body,
"hello",
"the second message content is restored"
);
// Deleting the conversation cascades to its messages (foreign_keys = ON).
await ConversationStore.deleteConversationById("conv-msgs");
const rows = await ConversationStore.connection.execute(
"SELECT count(*) AS c FROM message WHERE conv_id = :conv_id",
{ conv_id: "conv-msgs" }
);
Assert.equal(
rows[0].getResultByName("c"),
0,
"messages are cascade-deleted with the conversation"
);
});
add_task(async function test_removed_messages_are_pruned() {
const conversation = new Conversation({
id: "conv-prune",
createdDate: 1,
updatedDate: 1,
});
conversation.addMessage(MESSAGE_ROLE.USER, { type: "text", body: "one" }, 0);
conversation.addMessage(
MESSAGE_ROLE.ASSISTANT,
{ type: "text", body: "two" },
0
);
await ConversationStore.updateConversation(conversation);
// Drop the last message in memory (as retry/truncate would) and re-save.
conversation.messages = conversation.messages.slice(0, 1);
await ConversationStore.updateConversation(conversation);
const loaded = await ConversationStore.findConversationById("conv-prune");
Assert.equal(
loaded.messages.length,
1,
"the removed message is pruned, not resurrected"
);
Assert.equal(
loaded.messages[0].content.body,
"one",
"the remaining message is kept"
);
});
add_task(async function test_clearing_all_messages_removes_rows() {
const conversation = new Conversation({
id: "conv-clear",
createdDate: 1,
updatedDate: 1,
});
conversation.addMessage(MESSAGE_ROLE.USER, { type: "text", body: "one" }, 0);
conversation.addMessage(
MESSAGE_ROLE.ASSISTANT,
{ type: "text", body: "two" },
0
);
await ConversationStore.updateConversation(conversation);
// Emptying the conversation (as clearMessages would) must delete every row,
// exercising the empty keep-list branch of the delete.
conversation.messages = [];
await ConversationStore.updateConversation(conversation);
const rows = await ConversationStore.connection.execute(
"SELECT count(*) AS c FROM message WHERE conv_id = :conv_id",
{ conv_id: "conv-clear" }
);
Assert.equal(
rows[0].getResultByName("c"),
0,
"all messages are removed when the conversation is emptied"
);
});