mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-19 09:50:15 +01:00
ci: add unit tests for backend prompt file attachment code and for uploaded file processing
This commit is contained in:
Dustin Healy
parent
d8997fdf0e
commit
3fd84216cf
2 changed files with 1401 additions and 0 deletions
520
api/server/services/Files/process.spec.js
Normal file
520
api/server/services/Files/process.spec.js
Normal file
|
|
@ -0,0 +1,520 @@
|
|||
const { EToolResources, FileSources, FileContext } = require('librechat-data-provider');
|
||||
|
||||
jest.mock('~/server/services/Files/strategies', () => {
|
||||
const mockHandleFileUpload = jest.fn();
|
||||
const mockHandleImageUpload = jest.fn();
|
||||
|
||||
return {
|
||||
getStrategyFunctions: jest.fn((source) => ({
|
||||
handleFileUpload: mockHandleFileUpload.mockImplementation(({ file, file_id }) =>
|
||||
Promise.resolve({
|
||||
filepath: `/uploads/${source}/${file_id}`,
|
||||
bytes: file?.size || 20,
|
||||
}),
|
||||
),
|
||||
handleImageUpload: mockHandleImageUpload.mockImplementation(({ file, file_id }) =>
|
||||
Promise.resolve({
|
||||
filepath: `/uploads/${source}/images/${file_id}`,
|
||||
bytes: file.size,
|
||||
width: 800,
|
||||
height: 600,
|
||||
}),
|
||||
),
|
||||
})),
|
||||
};
|
||||
});
|
||||
|
||||
jest.mock('~/models/File', () => {
|
||||
const mockCreateFile = jest.fn();
|
||||
return {
|
||||
createFile: mockCreateFile.mockImplementation((fileInfo) =>
|
||||
Promise.resolve({ _id: 'test-file-id', ...fileInfo }),
|
||||
),
|
||||
updateFileUsage: jest.fn().mockResolvedValue(),
|
||||
};
|
||||
});
|
||||
|
||||
jest.mock('~/models/Agent', () => ({
|
||||
addAgentResourceFile: jest.fn().mockResolvedValue(),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/Config/getEndpointsConfig', () => ({
|
||||
checkCapability: jest.fn().mockResolvedValue(true),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/utils/getFileStrategy', () => ({
|
||||
getFileStrategy: jest.fn(() => {
|
||||
return 'local';
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/Files/VectorDB/crud', () => ({
|
||||
uploadVectors: jest.fn(({ file_id }) =>
|
||||
Promise.resolve({
|
||||
success: true,
|
||||
vectorIds: [`vector-${file_id}-1`, `vector-${file_id}-2`],
|
||||
}),
|
||||
),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/controllers/assistants/helpers', () => ({
|
||||
getOpenAIClient: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/Tools/credentials', () => ({
|
||||
loadAuthValues: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('fs', () => ({
|
||||
...jest.requireActual('fs'),
|
||||
createReadStream: jest.fn(() => 'mock-stream'),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/utils/queue', () => ({
|
||||
LB_QueueAsyncCall: jest.fn((fn, args, callback) => {
|
||||
if (callback) {
|
||||
callback(null, { success: true });
|
||||
}
|
||||
return Promise.resolve({ success: true });
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/Config/app', () => ({
|
||||
getAppConfig: jest.fn().mockResolvedValue({
|
||||
fileStrategy: 'local',
|
||||
fileStrategies: {
|
||||
agents: 'local',
|
||||
},
|
||||
imageOutputType: 'jpeg',
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/services/Files/images', () => ({
|
||||
processImageFile: jest.fn().mockResolvedValue({
|
||||
filepath: '/test/image/path',
|
||||
width: 800,
|
||||
height: 600,
|
||||
}),
|
||||
handleImageUpload: jest.fn().mockResolvedValue({
|
||||
filepath: '/test/image/uploaded/path',
|
||||
bytes: 1024,
|
||||
width: 800,
|
||||
height: 600,
|
||||
}),
|
||||
}));
|
||||
|
||||
describe('File Processing - processAgentFileUpload', () => {
|
||||
let processAgentFileUpload;
|
||||
let mockHandleFileUpload;
|
||||
let mockHandleImageUpload;
|
||||
let mockCreateFile;
|
||||
let mockAddAgentResourceFile;
|
||||
let mockUploadVectors;
|
||||
let mockCheckCapability;
|
||||
let mockGetFileStrategy;
|
||||
|
||||
beforeAll(() => {
|
||||
const processModule = require('./process');
|
||||
processAgentFileUpload = processModule.processAgentFileUpload;
|
||||
|
||||
const { getStrategyFunctions } = require('~/server/services/Files/strategies');
|
||||
const mockStrategies = getStrategyFunctions();
|
||||
mockHandleFileUpload = mockStrategies.handleFileUpload;
|
||||
mockHandleImageUpload = mockStrategies.handleImageUpload;
|
||||
|
||||
mockCreateFile = require('~/models/File').createFile;
|
||||
mockAddAgentResourceFile = require('~/models/Agent').addAgentResourceFile;
|
||||
mockUploadVectors = require('~/server/services/Files/VectorDB/crud').uploadVectors;
|
||||
mockCheckCapability = require('~/server/services/Config/getEndpointsConfig').checkCapability;
|
||||
mockGetFileStrategy = require('~/server/utils/getFileStrategy').getFileStrategy;
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('processAgentFileUpload', () => {
|
||||
it('should process image file upload for agent with proper file handling', async () => {
|
||||
const mockReq = {
|
||||
user: { id: 'test-user-id' },
|
||||
file: {
|
||||
buffer: Buffer.from('test image data'),
|
||||
mimetype: 'image/jpeg',
|
||||
size: 1024,
|
||||
originalname: 'test-image.jpg',
|
||||
},
|
||||
body: {
|
||||
file_id: 'test-file-id',
|
||||
},
|
||||
config: {
|
||||
fileStrategy: 'local',
|
||||
fileStrategies: {
|
||||
agents: 'local',
|
||||
},
|
||||
imageOutputType: 'jpeg',
|
||||
},
|
||||
};
|
||||
|
||||
const mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn(),
|
||||
};
|
||||
|
||||
const metadata = {
|
||||
agent_id: 'test-agent-id',
|
||||
tool_resource: EToolResources.image_edit,
|
||||
file_id: 'test-file-id',
|
||||
};
|
||||
|
||||
await processAgentFileUpload({ req: mockReq, res: mockRes, metadata });
|
||||
|
||||
expect(mockGetFileStrategy).toHaveBeenCalledWith(mockReq.config, { isImage: true });
|
||||
|
||||
expect(mockHandleImageUpload).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
req: mockReq,
|
||||
file: mockReq.file,
|
||||
file_id: expect.any(String),
|
||||
}),
|
||||
);
|
||||
|
||||
expect(mockCreateFile).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
user: 'test-user-id',
|
||||
file_id: 'test-file-id',
|
||||
bytes: 1024,
|
||||
filename: 'test-image.jpg',
|
||||
context: FileContext.agents,
|
||||
type: 'image/jpeg',
|
||||
source: FileSources.local,
|
||||
width: 800,
|
||||
height: 600,
|
||||
}),
|
||||
true,
|
||||
);
|
||||
|
||||
expect(mockAddAgentResourceFile).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
agent_id: 'test-agent-id',
|
||||
file_id: 'test-file-id',
|
||||
tool_resource: EToolResources.image_edit,
|
||||
req: mockReq,
|
||||
}),
|
||||
);
|
||||
|
||||
expect(mockRes.status).toHaveBeenCalledWith(200);
|
||||
expect(mockRes.json).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
message: 'Agent file uploaded and processed successfully',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should process file_search tool resource with dual storage (file + vector)', async () => {
|
||||
const mockReq = {
|
||||
user: { id: 'test-user-id' },
|
||||
file: {
|
||||
buffer: Buffer.from('test file data'),
|
||||
mimetype: 'application/pdf',
|
||||
size: 2048,
|
||||
originalname: 'test-document.pdf',
|
||||
},
|
||||
body: {
|
||||
file_id: 'test-file-id',
|
||||
},
|
||||
config: {
|
||||
fileStrategy: 'local',
|
||||
fileStrategies: {
|
||||
agents: 'local',
|
||||
},
|
||||
imageOutputType: 'jpeg',
|
||||
},
|
||||
};
|
||||
|
||||
const mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn(),
|
||||
};
|
||||
|
||||
const metadata = {
|
||||
agent_id: 'test-agent-id',
|
||||
tool_resource: EToolResources.file_search,
|
||||
file_id: 'test-file-id',
|
||||
};
|
||||
|
||||
await processAgentFileUpload({ req: mockReq, res: mockRes, metadata });
|
||||
|
||||
expect(mockGetFileStrategy).toHaveBeenCalledWith(mockReq.config, { isImage: false });
|
||||
|
||||
expect(mockHandleFileUpload).toHaveBeenCalledWith({
|
||||
req: mockReq,
|
||||
file: mockReq.file,
|
||||
file_id: 'test-file-id',
|
||||
basePath: 'uploads',
|
||||
entity_id: 'test-agent-id',
|
||||
});
|
||||
|
||||
expect(mockUploadVectors).toHaveBeenCalledWith({
|
||||
req: mockReq,
|
||||
file: mockReq.file,
|
||||
file_id: 'test-file-id',
|
||||
entity_id: 'test-agent-id',
|
||||
});
|
||||
|
||||
expect(mockCreateFile).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
user: 'test-user-id',
|
||||
file_id: 'test-file-id',
|
||||
filename: 'test-document.pdf',
|
||||
context: FileContext.agents,
|
||||
type: 'application/pdf',
|
||||
source: FileSources.local,
|
||||
bytes: 2048,
|
||||
filepath: '/uploads/local/test-file-id',
|
||||
metadata: {},
|
||||
}),
|
||||
true,
|
||||
);
|
||||
|
||||
expect(mockAddAgentResourceFile).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
agent_id: 'test-agent-id',
|
||||
file_id: 'test-file-id',
|
||||
tool_resource: EToolResources.file_search,
|
||||
req: mockReq,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle missing tool_resource parameter', async () => {
|
||||
const mockReq = {
|
||||
user: { id: 'test-user-id' },
|
||||
file: {
|
||||
buffer: Buffer.from('test file data'),
|
||||
mimetype: 'application/pdf',
|
||||
size: 2048,
|
||||
originalname: 'test-document.pdf',
|
||||
},
|
||||
body: {
|
||||
file_id: 'test-file-id',
|
||||
},
|
||||
config: {
|
||||
fileStrategy: 'local',
|
||||
fileStrategies: {
|
||||
agents: 'local',
|
||||
},
|
||||
imageOutputType: 'jpeg',
|
||||
},
|
||||
};
|
||||
|
||||
const mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn(),
|
||||
};
|
||||
|
||||
const metadata = {
|
||||
agent_id: 'test-agent-id',
|
||||
file_id: 'test-file-id',
|
||||
};
|
||||
|
||||
await expect(
|
||||
processAgentFileUpload({ req: mockReq, res: mockRes, metadata }),
|
||||
).rejects.toThrow('No tool resource provided for agent file upload');
|
||||
});
|
||||
|
||||
it('should handle missing agent_id parameter', async () => {
|
||||
const mockReq = {
|
||||
user: { id: 'test-user-id' },
|
||||
file: {
|
||||
buffer: Buffer.from('test file data'),
|
||||
mimetype: 'application/pdf',
|
||||
size: 2048,
|
||||
originalname: 'test-document.pdf',
|
||||
},
|
||||
body: {
|
||||
file_id: 'test-file-id',
|
||||
},
|
||||
config: {
|
||||
fileStrategy: 'local',
|
||||
fileStrategies: {
|
||||
agents: 'local',
|
||||
},
|
||||
imageOutputType: 'jpeg',
|
||||
},
|
||||
};
|
||||
|
||||
const mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn(),
|
||||
};
|
||||
|
||||
const metadata = {
|
||||
tool_resource: EToolResources.file_search,
|
||||
file_id: 'test-file-id',
|
||||
};
|
||||
|
||||
await expect(
|
||||
processAgentFileUpload({ req: mockReq, res: mockRes, metadata }),
|
||||
).rejects.toThrow('No agent ID provided for agent file upload');
|
||||
});
|
||||
|
||||
it('should handle image uploads for non-image tool resources', async () => {
|
||||
const mockReq = {
|
||||
user: { id: 'test-user-id' },
|
||||
file: {
|
||||
buffer: Buffer.from('test image data'),
|
||||
mimetype: 'image/jpeg',
|
||||
size: 1024,
|
||||
originalname: 'test-image.jpg',
|
||||
},
|
||||
body: {
|
||||
file_id: 'test-file-id',
|
||||
},
|
||||
config: {
|
||||
fileStrategy: 'local',
|
||||
fileStrategies: {
|
||||
agents: 'local',
|
||||
},
|
||||
imageOutputType: 'jpeg',
|
||||
},
|
||||
};
|
||||
|
||||
const mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn(),
|
||||
};
|
||||
|
||||
const metadata = {
|
||||
agent_id: 'test-agent-id',
|
||||
tool_resource: EToolResources.file_search,
|
||||
file_id: 'test-file-id',
|
||||
};
|
||||
|
||||
await expect(
|
||||
processAgentFileUpload({ req: mockReq, res: mockRes, metadata }),
|
||||
).rejects.toThrow('Image uploads are not supported for file search tool resources');
|
||||
});
|
||||
|
||||
it('should check execute_code capability and load auth values when processing code files', async () => {
|
||||
const mockReq = {
|
||||
user: { id: 'test-user-id' },
|
||||
file: {
|
||||
buffer: Buffer.from('print("hello world")'),
|
||||
mimetype: 'text/x-python',
|
||||
size: 20,
|
||||
originalname: 'test.py',
|
||||
path: '/tmp/test-file.py',
|
||||
},
|
||||
body: {
|
||||
file_id: 'test-file-id',
|
||||
},
|
||||
config: {
|
||||
fileStrategy: 'local',
|
||||
fileStrategies: {
|
||||
agents: 'local',
|
||||
},
|
||||
imageOutputType: 'jpeg',
|
||||
},
|
||||
};
|
||||
|
||||
const mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn(),
|
||||
};
|
||||
|
||||
const metadata = {
|
||||
agent_id: 'test-agent-id',
|
||||
tool_resource: EToolResources.execute_code,
|
||||
file_id: 'test-file-id',
|
||||
};
|
||||
|
||||
const mockLoadAuthValues = require('~/server/services/Tools/credentials').loadAuthValues;
|
||||
mockLoadAuthValues.mockResolvedValue({ CODE_API_KEY: 'test-key' });
|
||||
|
||||
await processAgentFileUpload({ req: mockReq, res: mockRes, metadata });
|
||||
|
||||
expect(mockCheckCapability).toHaveBeenCalledWith(mockReq, 'execute_code');
|
||||
|
||||
expect(mockLoadAuthValues).toHaveBeenCalledWith({
|
||||
userId: 'test-user-id',
|
||||
authFields: ['LIBRECHAT_CODE_API_KEY'],
|
||||
});
|
||||
|
||||
expect(mockHandleFileUpload).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
expect.objectContaining({
|
||||
req: mockReq,
|
||||
stream: 'mock-stream',
|
||||
filename: 'test.py',
|
||||
entity_id: 'test-agent-id',
|
||||
apiKey: undefined,
|
||||
}),
|
||||
);
|
||||
|
||||
expect(mockHandleFileUpload).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
expect.objectContaining({
|
||||
req: mockReq,
|
||||
file: mockReq.file,
|
||||
file_id: 'test-file-id',
|
||||
basePath: 'uploads',
|
||||
entity_id: 'test-agent-id',
|
||||
}),
|
||||
);
|
||||
|
||||
expect(mockAddAgentResourceFile).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
agent_id: 'test-agent-id',
|
||||
file_id: 'test-file-id',
|
||||
tool_resource: EToolResources.execute_code,
|
||||
req: mockReq,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw error when example capability (execute_code) is not enabled', async () => {
|
||||
const mockReq = {
|
||||
user: { id: 'test-user-id' },
|
||||
file: {
|
||||
buffer: Buffer.from('print("hello world")'),
|
||||
mimetype: 'text/x-python',
|
||||
size: 20,
|
||||
originalname: 'test.py',
|
||||
},
|
||||
body: {
|
||||
file_id: 'test-file-id',
|
||||
},
|
||||
config: {
|
||||
fileStrategy: 'local',
|
||||
fileStrategies: {
|
||||
agents: 'local',
|
||||
},
|
||||
imageOutputType: 'jpeg',
|
||||
},
|
||||
};
|
||||
|
||||
const mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn(),
|
||||
};
|
||||
|
||||
const metadata = {
|
||||
agent_id: 'test-agent-id',
|
||||
tool_resource: EToolResources.execute_code,
|
||||
file_id: 'test-file-id',
|
||||
};
|
||||
|
||||
mockCheckCapability.mockResolvedValueOnce(false);
|
||||
|
||||
await expect(
|
||||
processAgentFileUpload({ req: mockReq, res: mockRes, metadata }),
|
||||
).rejects.toThrow('Code execution is not enabled for Agents');
|
||||
|
||||
expect(mockCheckCapability).toHaveBeenCalledWith(mockReq, 'execute_code');
|
||||
|
||||
expect(mockHandleFileUpload).not.toHaveBeenCalled();
|
||||
expect(mockCreateFile).not.toHaveBeenCalled();
|
||||
expect(mockAddAgentResourceFile).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue