From 2cf910351a9edadfd64f498a002f5b448680395b Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Jul 2025 22:00:11 -0400 Subject: [PATCH] refactor: Enhance test setup and cleanup for file access control - Introduced modelsToCleanup array to track models added during tests for proper cleanup. - Updated afterAll hooks in test files to ensure all collections are cleared and only added models are deleted. - Improved consistency in model initialization across test files. - Added comments for clarity on cleanup processes and test data management. --- api/models/File.spec.js | 34 +++++++++++++++----- api/server/routes/files/files.agents.test.js | 22 ++++++++++++- api/server/routes/files/files.test.js | 23 +++++++++++-- 3 files changed, 67 insertions(+), 12 deletions(-) diff --git a/api/models/File.spec.js b/api/models/File.spec.js index 5eb39ce3d9..62388dbef6 100644 --- a/api/models/File.spec.js +++ b/api/models/File.spec.js @@ -12,6 +12,7 @@ let Agent; let AclEntry; let User; let AccessRole; +let modelsToCleanup = []; describe('File Access Control', () => { let mongoServer; @@ -22,23 +23,39 @@ describe('File Access Control', () => { await mongoose.connect(mongoUri); // Initialize all models - createModels(mongoose); + const models = createModels(mongoose); + + // Track which models we're adding + modelsToCleanup = Object.keys(models); // Register models on mongoose.models so methods can access them - const models = require('~/db/models'); - Object.assign(mongoose.models, models); + const dbModels = require('~/db/models'); + Object.assign(mongoose.models, dbModels); - File = models.File; - Agent = models.Agent; - AclEntry = models.AclEntry; - User = models.User; - AccessRole = models.AccessRole; + File = dbModels.File; + Agent = dbModels.Agent; + AclEntry = dbModels.AclEntry; + User = dbModels.User; + AccessRole = dbModels.AccessRole; // Seed default roles await seedDefaultRoles(); }); afterAll(async () => { + // Clean up all collections before disconnecting + const collections = mongoose.connection.collections; + for (const key in collections) { + await collections[key].deleteMany({}); + } + + // Clear only the models we added + for (const modelName of modelsToCleanup) { + if (mongoose.models[modelName]) { + delete mongoose.models[modelName]; + } + } + await mongoose.disconnect(); await mongoServer.stop(); }); @@ -48,6 +65,7 @@ describe('File Access Control', () => { await Agent.deleteMany({}); await AclEntry.deleteMany({}); await User.deleteMany({}); + // Don't delete AccessRole as they are seeded defaults needed for tests }); describe('hasAccessToFilesViaAgent', () => { diff --git a/api/server/routes/files/files.agents.test.js b/api/server/routes/files/files.agents.test.js index 4e6c5ca457..eebc591f8c 100644 --- a/api/server/routes/files/files.agents.test.js +++ b/api/server/routes/files/files.agents.test.js @@ -2,8 +2,9 @@ const express = require('express'); const request = require('supertest'); const mongoose = require('mongoose'); const { v4: uuidv4 } = require('uuid'); -const { createMethods } = require('@librechat/data-schemas'); +const { PERMISSION_BITS } = require('librechat-data-provider'); const { MongoMemoryServer } = require('mongodb-memory-server'); +const { createMethods } = require('@librechat/data-schemas'); const { createAgent } = require('~/models/Agent'); const { createFile } = require('~/models/File'); @@ -46,6 +47,7 @@ describe('File Routes - Agent Files Endpoint', () => { let AclEntry; // eslint-disable-next-line no-unused-vars let AccessRole; + let modelsToCleanup = []; beforeAll(async () => { mongoServer = await MongoMemoryServer.create(); @@ -56,6 +58,9 @@ describe('File Routes - Agent Files Endpoint', () => { const { createModels } = require('@librechat/data-schemas'); const models = createModels(mongoose); + // Track which models we're adding + modelsToCleanup = Object.keys(models); + // Register models on mongoose.models so methods can access them Object.assign(mongoose.models, models); @@ -86,15 +91,30 @@ describe('File Routes - Agent Files Endpoint', () => { }); afterAll(async () => { + // Clean up all collections before disconnecting + const collections = mongoose.connection.collections; + for (const key in collections) { + await collections[key].deleteMany({}); + } + + // Clear only the models we added + for (const modelName of modelsToCleanup) { + if (mongoose.models[modelName]) { + delete mongoose.models[modelName]; + } + } + await mongoose.disconnect(); await mongoServer.stop(); }); beforeEach(async () => { + // Clean up all test data await File.deleteMany({}); await Agent.deleteMany({}); await User.deleteMany({}); await AclEntry.deleteMany({}); + // Don't delete AccessRole as they are seeded defaults needed for tests // Create test users authorId = new mongoose.Types.ObjectId(); diff --git a/api/server/routes/files/files.test.js b/api/server/routes/files/files.test.js index 59aefa0d11..02aaf3c688 100644 --- a/api/server/routes/files/files.test.js +++ b/api/server/routes/files/files.test.js @@ -61,11 +61,11 @@ describe('File Routes - Delete with Agent Access', () => { let Agent; let AclEntry; let User; + let AccessRole; let methods; + let modelsToCleanup = []; // eslint-disable-next-line no-unused-vars let agentId; - // eslint-disable-next-line no-unused-vars - let AccessRole; beforeAll(async () => { mongoServer = await MongoMemoryServer.create(); @@ -76,6 +76,9 @@ describe('File Routes - Delete with Agent Access', () => { const { createModels } = require('@librechat/data-schemas'); const models = createModels(mongoose); + // Track which models we're adding + modelsToCleanup = Object.keys(models); + // Register models on mongoose.models so methods can access them Object.assign(mongoose.models, models); @@ -106,6 +109,19 @@ describe('File Routes - Delete with Agent Access', () => { }); afterAll(async () => { + // Clean up all collections before disconnecting + const collections = mongoose.connection.collections; + for (const key in collections) { + await collections[key].deleteMany({}); + } + + // Clear only the models we added + for (const modelName of modelsToCleanup) { + if (mongoose.models[modelName]) { + delete mongoose.models[modelName]; + } + } + await mongoose.disconnect(); await mongoServer.stop(); }); @@ -113,11 +129,12 @@ describe('File Routes - Delete with Agent Access', () => { beforeEach(async () => { jest.clearAllMocks(); - // Clear database + // Clear database - clean up all test data await File.deleteMany({}); await Agent.deleteMany({}); await User.deleteMany({}); await AclEntry.deleteMany({}); + // Don't delete AccessRole as they are seeded defaults needed for tests // Create test data authorId = new mongoose.Types.ObjectId();