From beaa8de6481c7d7d7e0f58d2d3db8879e05e22cb Mon Sep 17 00:00:00 2001
From: Yongtao Huang <yongtaoh2022@gmail.com>
Date: Mon, 8 Sep 2025 09:34:04 +0800
Subject: [PATCH] Fix: correct queryKey in useBatchUpdateDocMetadata and add
 test case (#25327)

---
 web/service/knowledge/use-metadata.spec.tsx | 84 +++++++++++++++++++++
 web/service/knowledge/use-metadata.ts       |  2 +-
 2 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 web/service/knowledge/use-metadata.spec.tsx

diff --git a/web/service/knowledge/use-metadata.spec.tsx b/web/service/knowledge/use-metadata.spec.tsx
new file mode 100644
index 000000000..3a11da726
--- /dev/null
+++ b/web/service/knowledge/use-metadata.spec.tsx
@@ -0,0 +1,84 @@
+import { DataType } from '@/app/components/datasets/metadata/types'
+import { act, renderHook } from '@testing-library/react'
+import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
+import { useBatchUpdateDocMetadata } from '@/service/knowledge/use-metadata'
+import { useDocumentListKey } from './use-document'
+
+// Mock the post function to avoid real network requests
+jest.mock('@/service/base', () => ({
+  post: jest.fn().mockResolvedValue({ success: true }),
+}))
+
+const NAME_SPACE = 'dataset-metadata'
+
+describe('useBatchUpdateDocMetadata', () => {
+  let queryClient: QueryClient
+
+  beforeEach(() => {
+    // Create a fresh QueryClient before each test
+    queryClient = new QueryClient()
+  })
+
+  // Wrapper for React Query context
+  const wrapper = ({ children }: { children: React.ReactNode }) => (
+    <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+  )
+
+  it('should correctly invalidate dataset and document caches', async () => {
+    const { result } = renderHook(() => useBatchUpdateDocMetadata(), { wrapper })
+
+    // Spy on queryClient.invalidateQueries
+    const invalidateSpy = jest.spyOn(queryClient, 'invalidateQueries')
+
+    // Correct payload type: each document has its own metadata_list array
+
+    const payload = {
+      dataset_id: 'dataset-1',
+      metadata_list: [
+        {
+          document_id: 'doc-1',
+          metadata_list: [
+            { key: 'title-1', id: '01', name: 'name-1', type: DataType.string, value: 'new title 01' },
+          ],
+        },
+        {
+          document_id: 'doc-2',
+          metadata_list: [
+            { key: 'title-2', id: '02', name: 'name-1', type: DataType.string, value: 'new title 02' },
+          ],
+        },
+      ],
+    }
+
+    // Execute the mutation
+    await act(async () => {
+      await result.current.mutateAsync(payload)
+    })
+
+    // Expect invalidateQueries to have been called exactly 5 times
+    expect(invalidateSpy).toHaveBeenCalledTimes(5)
+
+    // Dataset cache invalidation
+    expect(invalidateSpy).toHaveBeenNthCalledWith(1, {
+      queryKey: [NAME_SPACE, 'dataset', 'dataset-1'],
+    })
+
+    // Document list cache invalidation
+    expect(invalidateSpy).toHaveBeenNthCalledWith(2, {
+      queryKey: [NAME_SPACE, 'document', 'dataset-1'],
+    })
+
+    // useDocumentListKey cache invalidation
+    expect(invalidateSpy).toHaveBeenNthCalledWith(3, {
+      queryKey: [...useDocumentListKey, 'dataset-1'],
+    })
+
+    // Single document cache invalidation
+    expect(invalidateSpy.mock.calls.slice(3)).toEqual(
+      expect.arrayContaining([
+        [{ queryKey: [NAME_SPACE, 'document', 'dataset-1', 'doc-1'] }],
+        [{ queryKey: [NAME_SPACE, 'document', 'dataset-1', 'doc-2'] }],
+      ]),
+    )
+  })
+})
diff --git a/web/service/knowledge/use-metadata.ts b/web/service/knowledge/use-metadata.ts
index 5e9186f53..eb85142d9 100644
--- a/web/service/knowledge/use-metadata.ts
+++ b/web/service/knowledge/use-metadata.ts
@@ -119,7 +119,7 @@ export const useBatchUpdateDocMetadata = () => {
       })
       // meta data in document list
       await queryClient.invalidateQueries({
-        queryKey: [NAME_SPACE, 'dataset', payload.dataset_id],
+        queryKey: [NAME_SPACE, 'document', payload.dataset_id],
       })
       await queryClient.invalidateQueries({
         queryKey: [...useDocumentListKey, payload.dataset_id],