3 files changed, 69 insertions, 23 deletions
diff --git a/packages/e2e_tests/tests/api/assets.test.ts b/packages/e2e_tests/tests/api/assets.test.ts
index 78a5c7fe..accc5cf0 100644
--- a/packages/e2e_tests/tests/api/assets.test.ts
+++ b/packages/e2e_tests/tests/api/assets.test.ts
@@ -183,6 +183,7 @@ describe("Assets API", () => {
     expect(firstAsset).toEqual({
       id: uploadResponse1.assetId,
       assetType: "bannerImage",
+      fileName: "test.pdf",
     });
 
     // Attach second asset
@@ -205,6 +206,7 @@ describe("Assets API", () => {
     expect(secondAsset).toEqual({
       id: uploadResponse2.assetId,
       assetType: "bannerImage",
+      fileName: "test.pdf",
     });
 
     // Get bookmark and verify assets
@@ -221,8 +223,16 @@ describe("Assets API", () => {
 
     expect(bookmarkWithAssets?.assets).toEqual(
       expect.arrayContaining([
-        { id: uploadResponse1.assetId, assetType: "bannerImage" },
-        { id: uploadResponse2.assetId, assetType: "bannerImage" },
+        {
+          id: uploadResponse1.assetId,
+          assetType: "bannerImage",
+          fileName: "test.pdf",
+        },
+        {
+          id: uploadResponse2.assetId,
+          assetType: "bannerImage",
+          fileName: "test.pdf",
+        },
       ]),
     );
 
@@ -258,8 +268,16 @@ describe("Assets API", () => {
 
     expect(bookmarkAfterReplace?.assets).toEqual(
       expect.arrayContaining([
-        { id: uploadResponse3.assetId, assetType: "bannerImage" },
-        { id: uploadResponse2.assetId, assetType: "bannerImage" },
+        {
+          id: uploadResponse3.assetId,
+          assetType: "bannerImage",
+          fileName: "test.pdf",
+        },
+        {
+          id: uploadResponse2.assetId,
+          assetType: "bannerImage",
+          fileName: "test.pdf",
+        },
       ]),
     );
 
@@ -291,7 +309,11 @@ describe("Assets API", () => {
     );
 
     expect(bookmarkAfterDetach?.assets).toEqual([
-      { id: uploadResponse3.assetId, assetType: "bannerImage" },
+      {
+        id: uploadResponse3.assetId,
+        assetType: "bannerImage",
+        fileName: "test.pdf",
+      },
     ]);
   });
 });
diff --git a/packages/e2e_tests/tests/api/bookmarks.test.ts b/packages/e2e_tests/tests/api/bookmarks.test.ts
index e1cfc8d5..a6bc85e3 100644
--- a/packages/e2e_tests/tests/api/bookmarks.test.ts
+++ b/packages/e2e_tests/tests/api/bookmarks.test.ts
@@ -3,6 +3,7 @@ import { assert, beforeEach, describe, expect, inject, it } from "vitest";
 import { createKarakeepClient } from "@karakeep/sdk";
 
 import { createTestUser } from "../../utils/api";
+import { waitUntil } from "../../utils/general";
 
 describe("Bookmarks API", () => {
   const port = inject("karakeepPort");
@@ -353,9 +354,22 @@ describe("Bookmarks API", () => {
       },
     });
 
-    // Wait 3 seconds for the search index to be updated
-    // TODO: Replace with a check that all queues are empty
-    await new Promise((f) => setTimeout(f, 3000));
+    await waitUntil(async () => {
+      const { data, response, error } = await client.GET("/bookmarks/search", {
+        params: {
+          query: {
+            q: "test bookmark",
+          },
+        },
+      });
+      if (error) {
+        throw error;
+      }
+      if (response.status !== 200) {
+        throw new Error(`Search request failed with status ${response.status}`);
+      }
+      return (data?.bookmarks.length ?? 0) >= 2;
+    }, 'Search index contains the new bookmarks for query "test bookmark"');
 
     // Search for bookmarks
     const { data: searchResults, response: searchResponse } = await client.GET(
@@ -387,9 +401,23 @@ describe("Bookmarks API", () => {
 
     await Promise.all(bookmarkPromises);
 
-    // Wait 3 seconds for the search index to be updated
-    // TODO: Replace with a check that all queues are empty
-    await new Promise((f) => setTimeout(f, 3000));
+    await waitUntil(async () => {
+      const { data, response, error } = await client.GET("/bookmarks/search", {
+        params: {
+          query: {
+            q: "pagination",
+            limit: 5,
+          },
+        },
+      });
+      if (error) {
+        throw error;
+      }
+      if (response.status !== 200) {
+        throw new Error(`Search request failed with status ${response.status}`);
+      }
+      return (data?.bookmarks.length ?? 0) >= 5;
+    }, "Search index contains the pagination test bookmarks");
 
     // Get first page
     const { data: firstPage, response: firstResponse } = await client.GET(
diff --git a/packages/e2e_tests/tests/api/public.test.ts b/packages/e2e_tests/tests/api/public.test.ts
index 4ad37036..59f50495 100644
--- a/packages/e2e_tests/tests/api/public.test.ts
+++ b/packages/e2e_tests/tests/api/public.test.ts
@@ -90,18 +90,14 @@ describe("Public API", () => {
 
     const trpcClient = getTrpcClient(apiKey);
     // Wait for link bookmark to be crawled and have a banner image (screenshot)
-    await waitUntil(
-      async () => {
-        const res = await trpcClient.bookmarks.getBookmark.query({
-          bookmarkId: createBookmark1.id,
-        });
-        assert(res.content.type === BookmarkTypes.LINK);
-        // Check for screenshotAssetId as bannerImageUrl might be derived from it or original imageUrl
-        return !!res.content.screenshotAssetId || !!res.content.imageUrl;
-      },
-      "Bookmark is crawled and has banner info",
-      20000, // Increased timeout as crawling can take time
-    );
+    await waitUntil(async () => {
+      const res = await trpcClient.bookmarks.getBookmark.query({
+        bookmarkId: createBookmark1.id,
+      });
+      assert(res.content.type === BookmarkTypes.LINK);
+      // Check for screenshotAssetId as bannerImageUrl might be derived from it or original imageUrl
+      return !!res.content.screenshotAssetId || !!res.content.imageUrl;
+    }, "Bookmark is crawled and has banner info");
 
     const res = await trpcClient.publicBookmarks.getPublicBookmarksInList.query(
       {