| 1 | import base64 |
| 2 | from pathlib import Path |
| 3 | import re |
| 4 | import shutil |
| 5 | import subprocess |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | |
| 10 | PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| 11 | INPUT_STORE = PROJECT_ROOT / "webui/components/chat/input/input-store.js" |
| 12 | INDEX_JS = PROJECT_ROOT / "webui/index.js" |
| 13 | |
| 14 | |
| 15 | @pytest.mark.skipif(not shutil.which("node"), reason="Node.js is required") |
| 16 | def test_chat_input_keeps_separate_session_drafts() -> None: |
| 17 | index_source = INDEX_JS.read_text(encoding="utf-8") |
| 18 | set_context = index_source[index_source.index("export const setContext"):] |
| 19 | assert set_context.index("inputStore.setDraftContext(id);") < set_context.index("context = id;") |
| 20 | |
| 21 | source = INPUT_STORE.read_text(encoding="utf-8") |
| 22 | source = re.sub(r"^import .*?;\n", "", source, flags=re.MULTILINE) |
| 23 | source = source[: source.index('const store = createStore("chatInput", model);')] |
| 24 | module_source = r""" |
| 25 | const shortcuts = { |
| 26 | getCurrentContextId: () => globalThis.__context, |
| 27 | callJsonApi: async () => ({}), |
| 28 | frontendNotification: () => {}, |
| 29 | NotificationType: {}, |
| 30 | NotificationPriority: {}, |
| 31 | }; |
| 32 | const fileBrowserStore = {}; |
| 33 | const messageQueueStore = { hasQueue: false }; |
| 34 | const attachmentsStore = { |
| 35 | attachments: [], |
| 36 | clearAttachments() { this.attachments = []; }, |
| 37 | }; |
| 38 | const chatsStore = { selected: "", selectedContext: null }; |
| 39 | """ + source + "\nexport { model, chatsStore };\n" |
| 40 | module_url = "data:text/javascript;base64," + base64.b64encode( |
| 41 | module_source.encode("utf-8") |
| 42 | ).decode("ascii") |
| 43 | |
| 44 | script = f""" |
| 45 | const makeStorage = () => ({{ |
| 46 | values: new Map(), |
| 47 | getItem(key) {{ return this.values.get(key) ?? null; }}, |
| 48 | setItem(key, value) {{ this.values.set(key, String(value)); }}, |
| 49 | removeItem(key) {{ this.values.delete(key); }}, |
| 50 | }}); |
| 51 | globalThis.sessionStorage = makeStorage(); |
| 52 | globalThis.localStorage = makeStorage(); |
| 53 | globalThis.document = {{ activeElement: null, getElementById: () => null, querySelectorAll: () => [] }}; |
| 54 | globalThis.__context = null; |
| 55 | |
| 56 | const {{ model, chatsStore }} = await import({module_url!r}); |
| 57 | const assert = (condition, message) => {{ if (!condition) throw new Error(message); }}; |
| 58 | |
| 59 | globalThis.__context = "chat-a"; |
| 60 | model.setDraftContext("chat-a"); |
| 61 | model.message = "alpha draft"; |
| 62 | assert(sessionStorage.getItem("a0:chat-draft:chat-a") === "alpha draft", "chat A was not saved"); |
| 63 | |
| 64 | globalThis.__context = "chat-b"; |
| 65 | model.setDraftContext("chat-b"); |
| 66 | assert(model.message === "", "a new chat inherited another chat's draft"); |
| 67 | model.message = "beta draft"; |
| 68 | |
| 69 | globalThis.__context = "chat-a"; |
| 70 | model.setDraftContext("chat-a"); |
| 71 | assert(model.message === "alpha draft", "chat A was not restored"); |
| 72 | model.message = ""; |
| 73 | assert(sessionStorage.getItem("a0:chat-draft:chat-a") === null, "cleared draft remained stored"); |
| 74 | |
| 75 | globalThis.__context = null; |
| 76 | model.setDraftContext(""); |
| 77 | model.message = "welcome prompt"; |
| 78 | chatsStore.newChat = async () => {{ |
| 79 | globalThis.__context = "chat-new"; |
| 80 | chatsStore.selected = "chat-new"; |
| 81 | model.setDraftContext("chat-new"); |
| 82 | return "chat-new"; |
| 83 | }}; |
| 84 | let sent = ""; |
| 85 | globalThis.sendMessage = async () => {{ sent = model.message; }}; |
| 86 | await model.sendMessage(); |
| 87 | assert(sent === "welcome prompt", "creating a chat erased the Welcome prompt"); |
| 88 | assert(sessionStorage.getItem("a0:chat-draft:chat-new") === "welcome prompt", "first prompt did not follow its new chat"); |
| 89 | """ |
| 90 | |
| 91 | subprocess.run(["node", "--input-type=module", "-e", script], check=True, text=True) |