diff --git a/client/src/pages/__tests__/Workspace.test.tsx b/client/src/pages/__tests__/Workspace.test.tsx
index cd2694c..1605db0 100644
--- a/client/src/pages/__tests__/Workspace.test.tsx
+++ b/client/src/pages/__tests__/Workspace.test.tsx
@@ -601,15 +601,9 @@ describe("Workspace — proof gallery", () => {
     });
 
     await renderWorkspace();
-    await waitFor(
-      () => {
-        const gallery = screen.queryByTestId("proof-gallery");
-        if (gallery) {
-          expect(gallery).toBeInTheDocument();
-        }
-      },
-      { timeout: 2000 }
-    );
+    await waitFor(() => {
+      expect(screen.getByTestId("proof-gallery")).toBeInTheDocument();
+    });
   });
 
   it("shows no gallery panel when the selected feature has no proof", async () => {