Fix unit test for loading more data via requestSnapshot in the Electric collection

Author: kevin-dp

packages/electric-db-collection/src/electric.ts

@@ -15,12 +15,10 @@ import {
 import { compileSQL } from "./sql-compiler"
 import type {
   BaseCollectionConfig,
-  Collection,
   CollectionConfig,
   DeleteMutationFnParams,
   Fn,
   InsertMutationFnParams,
-  OnLoadMoreOptions,
   SyncConfig,
   UpdateMutationFnParams,
   UtilsRecord,
@@ -420,9 +418,9 @@ function createElectricSync<T extends Row<unknown>>(
       })
 
       return {
-        onLoadMore: (opts) => {
+        onLoadMore: async (opts) => {
           const snapshotParams = compileSQL<T>(opts)
-          return stream.requestSnapshot(snapshotParams)
+          await stream.requestSnapshot(snapshotParams)
         },
         cleanup: () => {
           // Unsubscribe from the stream
@@ -436,4 +434,3 @@ function createElectricSync<T extends Row<unknown>>(
     getSyncMetadata,
   }
 }
-

packages/electric-db-collection/tests/electric-live-query.test.ts

@@ -57,25 +57,34 @@ const mockSubscribe = vi.fn()
 const mockRequestSnapshot = vi.fn()
 const mockStream = {
   subscribe: mockSubscribe,
-  requestSnapshot: (...args: any) => {
-    mockRequestSnapshot(...args)
-    const results = mockRequestSnapshot.mock.results
-    const lastResult = results[results.length - 1]!.value
-
-    const subscribers = mockSubscribe.mock.calls.map(args => args[0])
-    subscribers.forEach(subscriber => subscriber(lastResult.data.map((row: any) => ({
-      type: `insert`,
+  requestSnapshot: async (...args: any) => {
+    const result = await mockRequestSnapshot(...args)
+    const subscribers = mockSubscribe.mock.calls.map((args) => args[0])
+    const data = [...result.data]
+
+    const messages: Array<Message<any>> = data.map((row: any) => ({
       value: row.value,
       key: row.key,
-    }))))
-  }
+      headers: row.headers,
+    }))
+
+    if (messages.length > 0) {
+      // add an up-to-date message
+      messages.push({
+        headers: { control: `up-to-date` },
+      })
+    }
+
+    subscribers.forEach((subscriber) => subscriber(messages))
+    return result
+  },
 }
 
 // Mock the requestSnapshot method
 // to return an empty array of data
 // since most tests don't use it
 mockRequestSnapshot.mockResolvedValue({
-  data: []
+  data: [],
 })
 
 vi.mock(`@electric-sql/client`, async () => {
@@ -458,14 +467,9 @@ describe.each([
     subscription.unsubscribe()
   })
   if (autoIndex === `eager`) {
-    it.only(`should load more data via requestSnapshot when creating live query with higher limit`, async () => {
-      // Reset mocks
-      vi.clearAllMocks()
+    it(`should load more data via requestSnapshot when creating live query with higher limit`, async () => {
       mockRequestSnapshot.mockResolvedValue({
-        data: [
-          { key: 5, value: { id: 5, name: `Eve`, age: 30, email: `[email protected]`, active: true } },
-          { key: 6, value: { id: 6, name: `Frank`, age: 35, email: `[email protected]`, active: true } },
-        ],
+        data: [],
       })
 
       // Initial sync with limited data
@@ -496,15 +500,45 @@ describe.each([
       expect(limitedLiveQuery.size).toBe(2) // Only first 2 active users
       expect(mockRequestSnapshot).toHaveBeenCalledTimes(1)
 
-      const callArgs = (index: number) => mockRequestSnapshot.mock.calls[index]?.[0]
+      const callArgs = (index: number) =>
+        mockRequestSnapshot.mock.calls[index]?.[0]
       expect(callArgs(0)).toMatchObject({
-        params: { "1": "true" },
-        where: "active = $1",
-        orderBy: "age NULLS FIRST",
+        params: { "1": `true` },
+        where: `active = $1`,
+        orderBy: `age NULLS FIRST`,
         limit: 2,
       })
 
-      // Create second live query with higher limit of 5
+      // Next call will return a snapshot containing 2 rows
+      // Calls after that will return the default empty snapshot
+      mockRequestSnapshot.mockResolvedValueOnce({
+        data: [
+          {
+            headers: { operation: `insert` },
+            key: 5,
+            value: {
+              id: 5,
+              name: `Eve`,
+              age: 30,
+              email: `[email protected]`,
+              active: true,
+            },
+          },
+          {
+            headers: { operation: `insert` },
+            key: 6,
+            value: {
+              id: 6,
+              name: `Frank`,
+              age: 35,
+              email: `[email protected]`,
+              active: true,
+            },
+          },
+        ],
+      })
+
+      // Create second live query with higher limit of 6
       const expandedLiveQuery = createLiveQueryCollection({
         id: `expanded-users-live-query`,
         startSync: true,
@@ -525,26 +559,34 @@ describe.each([
       await new Promise((resolve) => setTimeout(resolve, 0))
 
       // Verify that requestSnapshot was called with the correct parameters
-      expect(mockRequestSnapshot).toHaveBeenCalledTimes(3)
+      expect(mockRequestSnapshot).toHaveBeenCalledTimes(4)
 
       // Check that first it requested a limit of 6 users
       expect(callArgs(1)).toMatchObject({
-        params: { "1": "true" },
-        where: "active = $1",
-        orderBy: "age NULLS FIRST",
+        params: { "1": `true` },
+        where: `active = $1`,
+        orderBy: `age NULLS FIRST`,
         limit: 6,
       })
 
       // After this initial snapshot for the new live query it receives all 3 users from the local collection
       // so it still needs 3 more users to reach the limit of 6 so it requests 3 more to the sync layer
       expect(callArgs(2)).toMatchObject({
-        params: { "1": "true", "2": "25" },
-        where: "active = $1 AND age > $2",
-        orderBy: "age NULLS FIRST",
+        params: { "1": `true`, "2": `25` },
+        where: `active = $1 AND age > $2`,
+        orderBy: `age NULLS FIRST`,
         limit: 3,
       })
 
-      // The sync layer won't provide any more users so the DB is exhausted and it stops (i.e. doesn't request more) 
+      // The previous snapshot returned 2 more users so it still needs 1 more user to reach the limit of 6
+      expect(callArgs(3)).toMatchObject({
+        params: { "1": `true`, "2": `35` },
+        where: `active = $1 AND age > $2`,
+        orderBy: `age NULLS FIRST`,
+        limit: 1,
+      })
+
+      // The sync layer won't provide any more users so the DB is exhausted and it stops (i.e. doesn't request more)
 
       // The expanded live query should now have more data
       expect(expandedLiveQuery.status).toBe(`ready`)