From 28654199524b360beb821752e348539edd2d5830 Mon Sep 17 00:00:00 2001
From: "gru-agent[bot]" <185149714+gru-agent[bot]@users.noreply.github.com>
Date: Thu, 10 Jul 2025 18:28:39 +0800
Subject: [PATCH] Add unit tests for updateTrainingData API handler and modify
 dataId to be optional for retrying all error data. (#5194)

Co-authored-by: gru-agent[bot] <185149714+gru-agent[bot]@users.noreply.github.com>
---
 .../dataset/training/updateTrainingData.ts    |  23 ++-
 .../training/updateTrainingData.test.ts       | 143 ++++++++++++++++++
 2 files changed, 165 insertions(+), 1 deletion(-)
 create mode 100644 test/cases/pages/api/core/dataset/training/updateTrainingData.test.ts

diff --git a/projects/app/src/pages/api/core/dataset/training/updateTrainingData.ts b/projects/app/src/pages/api/core/dataset/training/updateTrainingData.ts
index f66433fd8..a8038d440 100644
--- a/projects/app/src/pages/api/core/dataset/training/updateTrainingData.ts
+++ b/projects/app/src/pages/api/core/dataset/training/updateTrainingData.ts
@@ -8,7 +8,7 @@ import { TrainingModeEnum } from '@fastgpt/global/core/dataset/constants';
 export type updateTrainingDataBody = {
   datasetId: string;
   collectionId: string;
-  dataId: string;
+  dataId?: string; // 改为可选，不传则重试所有错误数据
   q?: string;
   a?: string;
   chunkIndex?: number;
@@ -31,6 +31,25 @@ async function handler(
     per: WritePermissionVal
   });
 
+  // 如果没有传 dataId，则重试该集合下的所有错误数据
+  if (!dataId) {
+    await MongoDatasetTraining.updateMany(
+      {
+        teamId,
+        datasetId,
+        collectionId,
+        errorMsg: { $exists: true, $ne: null }
+      },
+      {
+        $unset: { errorMsg: '' },
+        retryCount: 3,
+        lockTime: new Date('2000')
+      }
+    );
+    return {};
+  }
+
+  // 单个数据重试逻辑
   const data = await MongoDatasetTraining.findOne({ teamId, datasetId, _id: dataId });
 
   if (!data) {
@@ -77,3 +96,5 @@ async function handler(
 }
 
 export default NextAPI(handler);
+
+export { handler };
diff --git a/test/cases/pages/api/core/dataset/training/updateTrainingData.test.ts b/test/cases/pages/api/core/dataset/training/updateTrainingData.test.ts
new file mode 100644
index 000000000..b55c6b662
--- /dev/null
+++ b/test/cases/pages/api/core/dataset/training/updateTrainingData.test.ts
@@ -0,0 +1,143 @@
+import { describe, expect, it, vi } from 'vitest';
+import { handler } from '@/pages/api/core/dataset/training/updateTrainingData';
+import { MongoDatasetTraining } from '@fastgpt/service/core/dataset/training/schema';
+import { authDatasetCollection } from '@fastgpt/service/support/permission/dataset/auth';
+import { TrainingModeEnum } from '@fastgpt/global/core/dataset/constants';
+
+vi.mock('@fastgpt/service/core/dataset/training/schema', () => ({
+  MongoDatasetTraining: {
+    findOne: vi.fn(),
+    updateOne: vi.fn(),
+    updateMany: vi.fn()
+  }
+}));
+
+vi.mock('@fastgpt/service/support/permission/dataset/auth', () => ({
+  authDatasetCollection: vi.fn()
+}));
+
+describe('updateTrainingData', () => {
+  it('should retry all error data when dataId is not provided', async () => {
+    vi.mocked(authDatasetCollection).mockResolvedValue({
+      teamId: 'team1'
+    });
+
+    const req = {
+      body: {
+        datasetId: 'dataset1',
+        collectionId: 'collection1'
+      }
+    };
+
+    await handler(req as any);
+
+    expect(MongoDatasetTraining.updateMany).toHaveBeenCalledWith(
+      {
+        teamId: 'team1',
+        datasetId: 'dataset1',
+        collectionId: 'collection1',
+        errorMsg: { $exists: true, $ne: null }
+      },
+      {
+        $unset: { errorMsg: '' },
+        retryCount: 3,
+        lockTime: new Date('2000')
+      }
+    );
+  });
+
+  it('should update single training data with image', async () => {
+    vi.mocked(authDatasetCollection).mockResolvedValue({
+      teamId: 'team1'
+    });
+
+    vi.mocked(MongoDatasetTraining.findOne).mockResolvedValue({
+      imageId: 'image1'
+    });
+
+    const req = {
+      body: {
+        datasetId: 'dataset1',
+        collectionId: 'collection1',
+        dataId: 'data1',
+        q: 'question',
+        a: 'answer',
+        chunkIndex: 1
+      }
+    };
+
+    await handler(req as any);
+
+    expect(MongoDatasetTraining.updateOne).toHaveBeenCalledWith(
+      {
+        teamId: 'team1',
+        datasetId: 'dataset1',
+        _id: 'data1'
+      },
+      {
+        $unset: { errorMsg: '' },
+        retryCount: 3,
+        mode: TrainingModeEnum.chunk,
+        q: 'question',
+        a: 'answer',
+        chunkIndex: 1,
+        lockTime: new Date('2000')
+      }
+    );
+  });
+
+  it('should update single training data without image', async () => {
+    vi.mocked(authDatasetCollection).mockResolvedValue({
+      teamId: 'team1'
+    });
+
+    vi.mocked(MongoDatasetTraining.findOne).mockResolvedValue({});
+
+    const req = {
+      body: {
+        datasetId: 'dataset1',
+        collectionId: 'collection1',
+        dataId: 'data1',
+        q: 'question',
+        a: 'answer',
+        chunkIndex: 1
+      }
+    };
+
+    await handler(req as any);
+
+    expect(MongoDatasetTraining.updateOne).toHaveBeenCalledWith(
+      {
+        teamId: 'team1',
+        datasetId: 'dataset1',
+        _id: 'data1'
+      },
+      {
+        $unset: { errorMsg: '' },
+        retryCount: 3,
+        q: 'question',
+        a: 'answer',
+        chunkIndex: 1,
+        lockTime: new Date('2000')
+      }
+    );
+  });
+
+  it('should reject when data not found', async () => {
+    vi.mocked(authDatasetCollection).mockResolvedValue({
+      teamId: 'team1'
+    });
+
+    vi.mocked(MongoDatasetTraining.findOne).mockResolvedValue(null);
+
+    const req = {
+      body: {
+        datasetId: 'dataset1',
+        collectionId: 'collection1',
+        dataId: 'data1'
+      }
+    };
+
+    await expect(handler(req as any)).rejects.toBe('data not found');
+  });
+});