@lobehub/chat 1.98.2 → 1.99.0

package/src/server/routers/lambda/generationBatch.test.ts

@@ -0,0 +1,376 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { GenerationBatchModel } from '@/database/models/generationBatch';
+import { GenerationBatchItem } from '@/database/schemas/generation';
+import { FileService } from '@/server/services/file';
+
+import { generationBatchRouter } from './generationBatch';
+
+vi.mock('@/database/models/generationBatch');
+vi.mock('@/server/services/file');
+
+describe('generationBatchRouter', () => {
+  const mockCtx = {
+    userId: 'test-user',
+    serverDB: {} as any,
+  };
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('should get generation batches by topic id', async () => {
+    const mockBatches = [
+      {
+        id: 'batch-1',
+        topicId: 'topic-1',
+        prompt: 'Test prompt',
+        generations: [
+          { id: 'gen-1', batchId: 'batch-1' },
+          { id: 'gen-2', batchId: 'batch-1' },
+        ],
+      },
+      {
+        id: 'batch-2',
+        topicId: 'topic-1',
+        prompt: 'Another prompt',
+        generations: [{ id: 'gen-3', batchId: 'batch-2' }],
+      },
+    ];
+
+    const mockQuery = vi.fn().mockResolvedValue(mockBatches);
+    vi.mocked(GenerationBatchModel).mockImplementation(
+      () =>
+        ({
+          queryGenerationBatchesByTopicIdWithGenerations: mockQuery,
+        }) as any,
+    );
+
+    const caller = generationBatchRouter.createCaller(mockCtx);
+
+    const result = await caller.getGenerationBatches({ topicId: 'topic-1' });
+
+    expect(result).toEqual(mockBatches);
+    expect(mockQuery).toHaveBeenCalledWith('topic-1');
+  });
+
+  it('should delete generation batch without thumbnails', async () => {
+    const mockBatchId = 'batch-123';
+    const mockDeletedBatch: GenerationBatchItem = {
+      id: mockBatchId,
+      userId: 'test-user',
+      generationTopicId: 'topic-1',
+      provider: 'test-provider',
+      model: 'test-model',
+      prompt: 'Test prompt',
+      width: 1024,
+      height: 1024,
+      ratio: null,
+      config: null,
+      accessedAt: new Date(),
+      createdAt: new Date(),
+      updatedAt: new Date(),
+    };
+
+    const mockDelete = vi.fn().mockResolvedValue({
+      deletedBatch: mockDeletedBatch,
+      thumbnailUrls: [], // 没有缩略图
+    });
+    const mockDeleteFiles = vi.fn();
+
+    vi.mocked(GenerationBatchModel).mockImplementation(
+      () =>
+        ({
+          delete: mockDelete,
+        }) as any,
+    );
+
+    vi.mocked(FileService).mockImplementation(
+      () =>
+        ({
+          deleteFiles: mockDeleteFiles,
+        }) as any,
+    );
+
+    const caller = generationBatchRouter.createCaller(mockCtx);
+    const result = await caller.deleteGenerationBatch({ batchId: mockBatchId });
+
+    expect(result).toEqual(mockDeletedBatch);
+    expect(mockDelete).toHaveBeenCalledWith(mockBatchId);
+    expect(mockDeleteFiles).not.toHaveBeenCalled(); // 没有文件要删除
+  });
+
+  it('should delete generation batch with thumbnails', async () => {
+    const mockBatchId = 'batch-123';
+    const mockThumbnailUrls = ['thumb1.jpg', 'thumb2.jpg', 'thumb3.jpg'];
+    const mockDeletedBatch: GenerationBatchItem = {
+      id: mockBatchId,
+      userId: 'test-user',
+      generationTopicId: 'topic-1',
+      provider: 'test-provider',
+      model: 'test-model',
+      prompt: 'Test prompt',
+      width: 1024,
+      height: 1024,
+      ratio: null,
+      config: null,
+      accessedAt: new Date(),
+      createdAt: new Date(),
+      updatedAt: new Date(),
+    };
+
+    const mockDelete = vi.fn().mockResolvedValue({
+      deletedBatch: mockDeletedBatch,
+      thumbnailUrls: mockThumbnailUrls,
+    });
+    const mockDeleteFiles = vi.fn().mockResolvedValue(true);
+
+    vi.mocked(GenerationBatchModel).mockImplementation(
+      () =>
+        ({
+          delete: mockDelete,
+        }) as any,
+    );
+
+    vi.mocked(FileService).mockImplementation(
+      () =>
+        ({
+          deleteFiles: mockDeleteFiles,
+        }) as any,
+    );
+
+    const caller = generationBatchRouter.createCaller(mockCtx);
+    const result = await caller.deleteGenerationBatch({ batchId: mockBatchId });
+
+    expect(result).toEqual(mockDeletedBatch);
+    expect(mockDelete).toHaveBeenCalledWith(mockBatchId);
+    expect(mockDeleteFiles).toHaveBeenCalledWith(mockThumbnailUrls);
+  });
+
+  it('should still return deleted batch when thumbnail deletion fails', async () => {
+    const mockBatchId = 'batch-123';
+    const mockThumbnailUrls = ['thumb1.jpg', 'thumb2.jpg'];
+    const mockDeletedBatch: GenerationBatchItem = {
+      id: mockBatchId,
+      userId: 'test-user',
+      generationTopicId: 'topic-1',
+      provider: 'test-provider',
+      model: 'test-model',
+      prompt: 'Test prompt',
+      width: 1024,
+      height: 1024,
+      ratio: null,
+      config: null,
+      accessedAt: new Date(),
+      createdAt: new Date(),
+      updatedAt: new Date(),
+    };
+
+    const mockDelete = vi.fn().mockResolvedValue({
+      deletedBatch: mockDeletedBatch,
+      thumbnailUrls: mockThumbnailUrls,
+    });
+
+    // Mock thumbnail deletion to fail
+    const mockDeleteFiles = vi.fn().mockRejectedValue(new Error('S3 thumbnail deletion failed'));
+    const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
+
+    vi.mocked(GenerationBatchModel).mockImplementation(
+      () =>
+        ({
+          delete: mockDelete,
+        }) as any,
+    );
+
+    vi.mocked(FileService).mockImplementation(
+      () =>
+        ({
+          deleteFiles: mockDeleteFiles,
+        }) as any,
+    );
+
+    const caller = generationBatchRouter.createCaller(mockCtx);
+    const result = await caller.deleteGenerationBatch({ batchId: mockBatchId });
+
+    // Database deletion should succeed even if thumbnail deletion fails
+    expect(result).toEqual(mockDeletedBatch);
+    expect(mockDelete).toHaveBeenCalledWith(mockBatchId);
+    expect(mockDeleteFiles).toHaveBeenCalledWith(mockThumbnailUrls);
+    expect(consoleSpy).toHaveBeenCalledWith(
+      'Failed to delete thumbnail files from S3:',
+      expect.any(Error),
+    );
+
+    consoleSpy.mockRestore();
+  });
+
+  it('should return undefined when deleting non-existent batch', async () => {
+    const mockBatchId = 'non-existent-batch';
+
+    const mockDelete = vi.fn().mockResolvedValue(undefined);
+    const mockDeleteFiles = vi.fn();
+
+    vi.mocked(GenerationBatchModel).mockImplementation(
+      () =>
+        ({
+          delete: mockDelete,
+        }) as any,
+    );
+
+    vi.mocked(FileService).mockImplementation(
+      () =>
+        ({
+          deleteFiles: mockDeleteFiles,
+        }) as any,
+    );
+
+    const caller = generationBatchRouter.createCaller(mockCtx);
+    const result = await caller.deleteGenerationBatch({ batchId: mockBatchId });
+
+    expect(result).toBeUndefined();
+    expect(mockDelete).toHaveBeenCalledWith(mockBatchId);
+    expect(mockDeleteFiles).not.toHaveBeenCalled(); // 没有文件要删除
+  });
+
+  it('should handle large number of thumbnails deletion', async () => {
+    const mockBatchId = 'batch-with-many-thumbnails';
+    // 模拟包含大量缩略图的批次
+    const mockThumbnailUrls = Array.from({ length: 50 }, (_, i) => `thumb${i + 1}.jpg`);
+    const mockDeletedBatch: GenerationBatchItem = {
+      id: mockBatchId,
+      userId: 'test-user',
+      generationTopicId: 'topic-1',
+      provider: 'test-provider',
+      model: 'test-model',
+      prompt: 'Batch with many generations',
+      width: 1024,
+      height: 1024,
+      ratio: null,
+      config: null,
+      accessedAt: new Date(),
+      createdAt: new Date(),
+      updatedAt: new Date(),
+    };
+
+    const mockDelete = vi.fn().mockResolvedValue({
+      deletedBatch: mockDeletedBatch,
+      thumbnailUrls: mockThumbnailUrls,
+    });
+    const mockDeleteFiles = vi.fn().mockResolvedValue(true);
+
+    vi.mocked(GenerationBatchModel).mockImplementation(
+      () =>
+        ({
+          delete: mockDelete,
+        }) as any,
+    );
+
+    vi.mocked(FileService).mockImplementation(
+      () =>
+        ({
+          deleteFiles: mockDeleteFiles,
+        }) as any,
+    );
+
+    const caller = generationBatchRouter.createCaller(mockCtx);
+    const result = await caller.deleteGenerationBatch({ batchId: mockBatchId });
+
+    expect(result).toEqual(mockDeletedBatch);
+    expect(mockDelete).toHaveBeenCalledWith(mockBatchId);
+    expect(mockDeleteFiles).toHaveBeenCalledWith(mockThumbnailUrls);
+    expect(mockDeleteFiles).toHaveBeenCalledTimes(1);
+  });
+
+  it('should handle empty generation batches result', async () => {
+    const mockQuery = vi.fn().mockResolvedValue([]);
+    vi.mocked(GenerationBatchModel).mockImplementation(
+      () =>
+        ({
+          queryGenerationBatchesByTopicIdWithGenerations: mockQuery,
+        }) as any,
+    );
+
+    const caller = generationBatchRouter.createCaller(mockCtx);
+
+    const result = await caller.getGenerationBatches({ topicId: 'non-existent-topic' });
+
+    expect(result).toEqual([]);
+    expect(mockQuery).toHaveBeenCalledWith('non-existent-topic');
+  });
+
+  it('should handle query error gracefully', async () => {
+    const mockQuery = vi.fn().mockRejectedValue(new Error('Database connection failed'));
+    vi.mocked(GenerationBatchModel).mockImplementation(
+      () =>
+        ({
+          queryGenerationBatchesByTopicIdWithGenerations: mockQuery,
+        }) as any,
+    );
+
+    const caller = generationBatchRouter.createCaller(mockCtx);
+
+    await expect(caller.getGenerationBatches({ topicId: 'topic-1' })).rejects.toThrow(
+      'Database connection failed',
+    );
+    expect(mockQuery).toHaveBeenCalledWith('topic-1');
+  });
+
+  it('should handle partial thumbnail deletion failure gracefully', async () => {
+    const mockBatchId = 'batch-123';
+    const mockThumbnailUrls = ['thumb1.jpg', 'thumb2.jpg', 'thumb3.jpg'];
+    const mockDeletedBatch: GenerationBatchItem = {
+      id: mockBatchId,
+      userId: 'test-user',
+      generationTopicId: 'topic-1',
+      provider: 'test-provider',
+      model: 'test-model',
+      prompt: 'Test prompt',
+      width: 1024,
+      height: 1024,
+      ratio: null,
+      config: null,
+      accessedAt: new Date(),
+      createdAt: new Date(),
+      updatedAt: new Date(),
+    };
+
+    const mockDelete = vi.fn().mockResolvedValue({
+      deletedBatch: mockDeletedBatch,
+      thumbnailUrls: mockThumbnailUrls,
+    });
+
+    // Mock partial failure - some thumbnails could not be deleted
+    const mockDeleteFiles = vi
+      .fn()
+      .mockRejectedValue(new Error('Some thumbnails could not be deleted from S3'));
+    const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
+
+    vi.mocked(GenerationBatchModel).mockImplementation(
+      () =>
+        ({
+          delete: mockDelete,
+        }) as any,
+    );
+
+    vi.mocked(FileService).mockImplementation(
+      () =>
+        ({
+          deleteFiles: mockDeleteFiles,
+        }) as any,
+    );
+
+    const caller = generationBatchRouter.createCaller(mockCtx);
+    const result = await caller.deleteGenerationBatch({ batchId: mockBatchId });
+
+    // Even with partial thumbnail deletion failure, batch deletion should succeed
+    expect(result).toEqual(mockDeletedBatch);
+    expect(mockDelete).toHaveBeenCalledWith(mockBatchId);
+    expect(mockDeleteFiles).toHaveBeenCalledWith(mockThumbnailUrls);
+    expect(consoleSpy).toHaveBeenCalledWith(
+      'Failed to delete thumbnail files from S3:',
+      expect.any(Error),
+    );
+
+    consoleSpy.mockRestore();
+  });
+});

package/src/server/routers/lambda/generationBatch.ts

@@ -0,0 +1,56 @@
+import { z } from 'zod';
+
+import { GenerationBatchModel } from '@/database/models/generationBatch';
+import { authedProcedure, router } from '@/libs/trpc/lambda';
+import { serverDatabase } from '@/libs/trpc/lambda/middleware';
+import { FileService } from '@/server/services/file';
+
+const generationBatchProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
+  const { ctx } = opts;
+
+  return opts.next({
+    ctx: {
+      fileService: new FileService(ctx.serverDB, ctx.userId),
+      generationBatchModel: new GenerationBatchModel(ctx.serverDB, ctx.userId),
+    },
+  });
+});
+
+export const generationBatchRouter = router({
+  deleteGenerationBatch: generationBatchProcedure
+    .input(z.object({ batchId: z.string() }))
+    .mutation(async ({ ctx, input }) => {
+      // 1. Delete database records and get thumbnail URLs to clean
+      const result = await ctx.generationBatchModel.delete(input.batchId);
+
+      // If batch not found, return early
+      if (!result) {
+        return;
+      }
+
+      const { deletedBatch, thumbnailUrls } = result;
+
+      // 2. Clean up thumbnail files from S3
+      // Note: Even if file deletion fails, we consider the batch deletion successful
+      // since the database record has been removed and users won't see the batch anymore
+      if (thumbnailUrls.length > 0) {
+        try {
+          await ctx.fileService.deleteFiles(thumbnailUrls);
+        } catch (error) {
+          // Log the error but don't throw - file cleanup failure shouldn't affect
+          // the user experience since the database operation succeeded
+          console.error('Failed to delete thumbnail files from S3:', error);
+        }
+      }
+
+      return deletedBatch;
+    }),
+
+  getGenerationBatches: generationBatchProcedure
+    .input(z.object({ topicId: z.string() }))
+    .query(async ({ ctx, input }) => {
+      return ctx.generationBatchModel.queryGenerationBatchesByTopicIdWithGenerations(input.topicId);
+    }),
+});
+
+export type GenerationBatchRouter = typeof generationBatchRouter;