tests: extract onAfterSend + implement sent-copy.test.ts

src/background/background.ts

@@ -23,6 +23,7 @@ import { toEmail, EMAIL_ATTRIBUTE_TYPE, findHtmlBody } from "../lib/utils";
 import { getOrCreateLocalFolder } from "../lib/folders";
 import { isPGEncrypted, wasPGEncrypted } from "../lib/detection";
 import { dispatchRuntimeMessage } from "./runtime-router";
+import { handleAfterSend } from "./sent-copy";
 import type {
   Policy,
   SerializedRecipient,
@@ -102,37 +103,13 @@ browser.runtime.onMessage.addListener(
 
 browser.compose.onBeforeSend.addListener(handleBeforeSend);
 
-browser.compose.onAfterSend.addListener(async (tab, sendInfo) => {
-  const state = composeTabs.get(tab.id);
-  if (!state?.sentMimeData) return;
-
-  try {
-    for (const msg of sendInfo.messages) {
-      if (await isPGEncrypted(msg.id)) {
-        const localFolder = await getOrCreateLocalFolder("PostGuard Sent");
-        if (localFolder) {
-          const file = new File([state.sentMimeData as BlobPart], "sent.eml", {
-            type: "text/plain",
-          });
-          const localMsg = await browser.messages.import(
-            file,
-            localFolder.id
-          );
-          await browser.messages.move([localMsg.id], msg.folder.id);
-          await browser.messages.delete([msg.id], true);
-        }
-      }
-    }
-  } catch (e) {
-    console.error("[PostGuard] Failed to manage sent copy:", e);
-    notifyError("sentCopyError");
-  } finally {
-    cleanupComposeTab(tab.id);
-    clearInFlightUpload(tab.id);
-    persistEncryptState().catch(console.warn);
-    persistInFlightUploads().catch(console.warn);
-  }
-});
+browser.compose.onAfterSend.addListener((tab, sendInfo) =>
+  handleAfterSend(tab, sendInfo, {
+    notifyError,
+    isPGEncrypted,
+    getOrCreateLocalFolder,
+  }),
+);
 
 // Clean up decryptedMessages when messages are deleted
 browser.messages.onDeleted.addListener((deletedMessages) => {

src/background/sent-copy.ts

@@ -0,0 +1,71 @@
+/// <reference path="../types/thunderbird.d.ts" />
+
+import {
+  composeTabs,
+  cleanupComposeTab,
+  clearInFlightUpload,
+  persistEncryptState,
+  persistInFlightUploads,
+} from "./state";
+
+// Pulled out of background.ts so the post-send cleanup can be exercised
+// under vitest without standing up the whole background script. Deps that
+// hit Thunderbird APIs we cannot easily mock from the browser-mock helper
+// are injected, and the rest goes through the browser-mock surface.
+
+export interface SentMessage {
+  id: number;
+  folder: { id: unknown };
+}
+
+export interface SentInfo {
+  messages: SentMessage[];
+}
+
+export interface SentCopyDeps {
+  notifyError: (messageKey: string) => void;
+  isPGEncrypted: (msgId: number) => Promise<boolean>;
+  getOrCreateLocalFolder: (
+    name: string,
+  ) => Promise<{ id: unknown } | undefined>;
+}
+
+/**
+ * Post-send handler. For each message the user just sent, if it is a
+ * PostGuard ciphertext, swap the encrypted copy in Sent with the plaintext
+ * MIME stashed during encryption. State cleanup happens in `finally` so a
+ * failure mid-swap never strands the compose tab in `composeTabs`.
+ */
+export async function handleAfterSend(
+  tab: { id: number },
+  sendInfo: SentInfo,
+  deps: SentCopyDeps,
+): Promise<void> {
+  const state = composeTabs.get(tab.id);
+  if (!state?.sentMimeData) return;
+
+  try {
+    for (const msg of sendInfo.messages) {
+      if (!(await deps.isPGEncrypted(msg.id))) continue;
+      const localFolder = await deps.getOrCreateLocalFolder("PostGuard Sent");
+      if (!localFolder) continue;
+      const file = new File([state.sentMimeData as BlobPart], "sent.eml", {
+        type: "text/plain",
+      });
+      const localMsg = await browser.messages.import(
+        file,
+        localFolder.id as any,
+      );
+      await browser.messages.move([localMsg.id], msg.folder.id as any);
+      await browser.messages.delete([msg.id], true);
+    }
+  } catch (e) {
+    console.error("[PostGuard] Failed to manage sent copy:", e);
+    deps.notifyError("sentCopyError");
+  } finally {
+    cleanupComposeTab(tab.id);
+    clearInFlightUpload(tab.id);
+    persistEncryptState().catch(console.warn);
+    persistInFlightUploads().catch(console.warn);
+  }
+}

tests/sent-copy.test.ts

@@ -1,22 +1,141 @@
-import { describe, it, expect } from "vitest";
+import { describe, it, expect, beforeEach, vi } from "vitest";
+import { handleAfterSend, type SentInfo } from "../src/background/sent-copy";
+import {
+  composeTabs,
+  decryptedMessages,
+  inFlightUploads,
+  pendingCryptoPopups,
+  pendingPolicyEditors,
+} from "../src/background/state";
+import { installBrowserMock, type BrowserMock } from "./helpers/browser-mock";
 
-// These tests verify the onAfterSend sent copy management.
-// They require browser API mocks.
+let mock: BrowserMock;
+let notifyError: ReturnType<typeof vi.fn>;
+let isPGEncrypted: ReturnType<typeof vi.fn>;
+let getOrCreateLocalFolder: ReturnType<typeof vi.fn>;
+let importSpy: ReturnType<typeof vi.spyOn>;
+let moveSpy: ReturnType<typeof vi.spyOn>;
+let deleteSpy: ReturnType<typeof vi.spyOn>;
+
+const SENT_FOLDER = { id: "sent://Sent" };
+const LOCAL_FOLDER = { id: "local://PostGuard Sent" };
+
+function deps() {
+  return { notifyError, isPGEncrypted, getOrCreateLocalFolder };
+}
+
+function sentInfo(messageIds: number[] = [1]): SentInfo {
+  return {
+    messages: messageIds.map((id) => ({ id, folder: SENT_FOLDER })),
+  };
+}
+
+beforeEach(() => {
+  mock = installBrowserMock();
+  composeTabs.clear();
+  decryptedMessages.clear();
+  pendingCryptoPopups.clear();
+  pendingPolicyEditors.clear();
+  inFlightUploads.clear();
+
+  notifyError = vi.fn();
+  isPGEncrypted = vi.fn(async () => true);
+  getOrCreateLocalFolder = vi.fn(async () => LOCAL_FOLDER);
+
+  importSpy = vi.spyOn(browser.messages, "import");
+  moveSpy = vi.spyOn(browser.messages, "move");
+  deleteSpy = vi.spyOn(browser.messages, "delete");
+});
 
 describe("onAfterSend — sent copy management", () => {
-  it.todo("should skip when no sentMimeData is stored");
+  it("should skip when no sentMimeData is stored", async () => {
+    composeTabs.set(7, { encrypt: true });
+    await handleAfterSend({ id: 7 }, sentInfo([1]), deps());
+
+    expect(isPGEncrypted).not.toHaveBeenCalled();
+    expect(importSpy).not.toHaveBeenCalled();
+  });
+
+  it("should skip non-PG-encrypted sent messages", async () => {
+    composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+    isPGEncrypted.mockResolvedValueOnce(false);
+
+    await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
+
+    expect(isPGEncrypted).toHaveBeenCalledWith(42);
+    expect(getOrCreateLocalFolder).not.toHaveBeenCalled();
+    expect(importSpy).not.toHaveBeenCalled();
+  });
+
+  it("should import plaintext MIME into PostGuard Sent folder", async () => {
+    composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+
+    await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
+
+    expect(getOrCreateLocalFolder).toHaveBeenCalledWith("PostGuard Sent");
+    expect(importSpy).toHaveBeenCalledTimes(1);
+    const [file, folder] = importSpy.mock.calls[0] as [File, unknown];
+    expect(file).toBeInstanceOf(File);
+    expect(file.name).toBe("sent.eml");
+    expect(folder).toBe(LOCAL_FOLDER.id);
+  });
+
+  it("should move imported message to original sent folder", async () => {
+    composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+
+    await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
+
+    expect(moveSpy).toHaveBeenCalledWith([9999], SENT_FOLDER.id);
+  });
+
+  it("should delete the encrypted copy from sent folder", async () => {
+    composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+
+    await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
+
+    expect(deleteSpy).toHaveBeenCalledWith([42], true);
+  });
+
+  it("should notify user when sent copy management fails", async () => {
+    composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+    importSpy.mockRejectedValueOnce(new Error("import boom"));
+
+    await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
+
+    expect(notifyError).toHaveBeenCalledWith("sentCopyError");
+  });
+
+  it("should clean up compose tab state after send", async () => {
+    composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+    inFlightUploads.set(7, { uuid: "u", recoveryToken: "r", startedAt: Date.now() });
+
+    await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
 
-  it.todo("should skip non-PG-encrypted sent messages");
+    expect(composeTabs.has(7)).toBe(false);
+    expect(inFlightUploads.has(7)).toBe(false);
+  });
 
-  it.todo("should import plaintext MIME into PostGuard Sent folder");
+  it("should clean up compose tab state even on failure", async () => {
+    composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+    inFlightUploads.set(7, { uuid: "u", recoveryToken: "r", startedAt: Date.now() });
+    importSpy.mockRejectedValueOnce(new Error("nope"));
 
-  it.todo("should move imported message to original sent folder");
+    await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
 
-  it.todo("should delete the encrypted copy from sent folder");
+    expect(composeTabs.has(7)).toBe(false);
+    expect(inFlightUploads.has(7)).toBe(false);
+  });
 
-  it.todo("should notify user when sent copy management fails");
+  it("should skip the swap when no local folder is available", async () => {
+    composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) });
+    getOrCreateLocalFolder.mockResolvedValueOnce(undefined);
 
-  it.todo("should clean up compose tab state after send");
+    await handleAfterSend({ id: 7 }, sentInfo([42]), deps());
 
-  it.todo("should clean up compose tab state even on failure");
+    expect(importSpy).not.toHaveBeenCalled();
+    expect(moveSpy).not.toHaveBeenCalled();
+    expect(deleteSpy).not.toHaveBeenCalled();
+    // Cleanup still runs.
+    expect(composeTabs.has(7)).toBe(false);
+  });
 });