test: fix flaky extension tests with polling-based waits

Replace fixed setTimeout delays with a polling mechanism that waits for
expected state changes. This resolves intermittent test failures in CI
where 500ms delays were insufficient for VS Code operations to complete.

- Add waitForCondition helper to poll for state changes with timeout
- Update toggleAudited test to poll until audit state changes
- Update addPartiallyAudited test to poll for state changes
- Update showMarkedFilesDayLog test to poll for editor visibility

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: fegge

test/extension/suite/commands.test.ts

@@ -32,6 +32,21 @@ function readWeauditData(workspaceFolder: vscode.WorkspaceFolder): SerializedDat
     return JSON.parse(content) as SerializedData;
 }
 
+/**
+ * Poll for a condition to become true, with timeout.
+ * Returns true if condition was met, false if timed out.
+ */
+async function waitForCondition(condition: () => boolean, timeoutMs: number = 5000, pollIntervalMs: number = 100): Promise<boolean> {
+    const startTime = Date.now();
+    while (Date.now() - startTime < timeoutMs) {
+        if (condition()) {
+            return true;
+        }
+        await new Promise((resolve) => setTimeout(resolve, pollIntervalMs));
+    }
+    return condition(); // Final check
+}
+
 suite("Command Execution", () => {
     const extensionId = "trailofbits.weaudit";
     let testFileUri: vscode.Uri;
@@ -207,21 +222,26 @@ suite("Command Execution", () => {
         // If not audited, mark it as audited
         if (!wasAudited) {
             await vscode.commands.executeCommand("weAudit.toggleAudited");
-            await new Promise((resolve) => setTimeout(resolve, 500));
+            // Wait for the audit to complete
+            await waitForCondition(() => {
+                const data = readWeauditData(workspaceFolder);
+                return data?.auditedFiles.some((f) => f.path === relativePath) ?? false;
+            });
         }
 
         // Close all editors first to get a clean slate
         await vscode.commands.executeCommand("workbench.action.closeAllEditors");
-        await new Promise((resolve) => setTimeout(resolve, 200));
+        // Wait for editors to close
+        await waitForCondition(() => vscode.window.visibleTextEditors.length === 0, 2000);
 
         const editorsBefore = vscode.window.visibleTextEditors.length;
 
         // Now call the day log command - it should open a markdown document
         await vscode.commands.executeCommand("weAudit.showMarkedFilesDayLog");
-        await new Promise((resolve) => setTimeout(resolve, 500));
 
-        const editorsAfter = vscode.window.visibleTextEditors.length;
-        assert.ok(editorsAfter > editorsBefore, "A new document should be opened with the day log");
+        // Wait for a new editor to appear (poll instead of fixed delay)
+        const editorOpened = await waitForCondition(() => vscode.window.visibleTextEditors.length > editorsBefore, 5000);
+        assert.ok(editorOpened, "A new document should be opened with the day log");
 
         // Verify the opened document contains expected content
         const activeEditor = vscode.window.activeTextEditor;
@@ -235,7 +255,11 @@ suite("Command Execution", () => {
             await vscode.commands.executeCommand("workbench.action.closeActiveEditor");
             await openTestFile();
             await vscode.commands.executeCommand("weAudit.toggleAudited");
-            await new Promise((resolve) => setTimeout(resolve, 500));
+            // Wait for the unaudit to complete
+            await waitForCondition(() => {
+                const data = readWeauditData(workspaceFolder);
+                return !data?.auditedFiles.some((f) => f.path === relativePath);
+            });
         }
     });
 });

test/extension/suite/decorations.test.ts

@@ -25,6 +25,21 @@ function readWeauditData(workspaceFolder: vscode.WorkspaceFolder): SerializedDat
     return JSON.parse(content) as SerializedData;
 }
 
+/**
+ * Poll for a condition to become true, with timeout.
+ * Returns true if condition was met, false if timed out.
+ */
+async function waitForCondition(condition: () => boolean, timeoutMs: number = 5000, pollIntervalMs: number = 100): Promise<boolean> {
+    const startTime = Date.now();
+    while (Date.now() - startTime < timeoutMs) {
+        if (condition()) {
+            return true;
+        }
+        await new Promise((resolve) => setTimeout(resolve, pollIntervalMs));
+    }
+    return condition(); // Final check
+}
+
 suite("Editor Decorations", () => {
     const extensionId = "trailofbits.weaudit";
     let testFileUri: vscode.Uri;
@@ -80,13 +95,15 @@ suite("Editor Decorations", () => {
 
         // Toggle audited status
         await vscode.commands.executeCommand("weAudit.toggleAudited");
-        await new Promise((resolve) => setTimeout(resolve, 500));
 
-        // Verify the state changed
-        const dataAfter = readWeauditData(workspaceFolder);
-        const isAudited = dataAfter?.auditedFiles.some((f) => f.path === relativePath) ?? false;
+        // Wait for the state to actually change (poll instead of fixed delay)
+        const stateChanged = await waitForCondition(() => {
+            const dataAfter = readWeauditData(workspaceFolder);
+            const isAudited = dataAfter?.auditedFiles.some((f) => f.path === relativePath) ?? false;
+            return isAudited !== wasAudited;
+        });
 
-        assert.notStrictEqual(isAudited, wasAudited, "Audited status should be toggled");
+        assert.ok(stateChanged, "Audited status should be toggled");
     });
 
     test("addPartiallyAudited command adds partially audited region", async function () {
@@ -100,7 +117,11 @@ suite("Editor Decorations", () => {
         const dataCheck = readWeauditData(workspaceFolder);
         if (dataCheck?.auditedFiles.some((f) => f.path === relativePath)) {
             await vscode.commands.executeCommand("weAudit.toggleAudited");
-            await new Promise((resolve) => setTimeout(resolve, 500));
+            // Wait for the unaudit to complete
+            await waitForCondition(() => {
+                const data = readWeauditData(workspaceFolder);
+                return !data?.auditedFiles.some((f) => f.path === relativePath);
+            });
         }
 
         // Use lines that are not already partially audited
@@ -119,15 +140,17 @@ suite("Editor Decorations", () => {
 
         // Add partially audited region
         await vscode.commands.executeCommand("weAudit.addPartiallyAudited");
-        await new Promise((resolve) => setTimeout(resolve, 500));
 
-        // Verify state changed
-        const dataAfter = readWeauditData(workspaceFolder);
-        const regionExistsAfter =
-            dataAfter?.partiallyAuditedFiles.some((f) => f.path === relativePath && f.startLine === startLine && f.endLine === endLine) ?? false;
+        // Wait for the state to actually change (poll instead of fixed delay)
+        const stateChanged = await waitForCondition(() => {
+            const dataAfter = readWeauditData(workspaceFolder);
+            const regionExistsAfter =
+                dataAfter?.partiallyAuditedFiles.some((f) => f.path === relativePath && f.startLine === startLine && f.endLine === endLine) ?? false;
+            return regionExistsAfter !== regionExistsBefore;
+        });
 
         // The command toggles the region - so if it existed before it's removed, if it didn't exist it's added
-        assert.notStrictEqual(regionExistsAfter, regionExistsBefore, "Partially audited region should be toggled");
+        assert.ok(stateChanged, "Partially audited region should be toggled");
     });
 
     test("Extension survives file edits", async function () {