Upg: tests now uses resources (#10427)

* Upg: tests now uses resources

* Fix: csv tests

front/lib/api/data_sources.ts

@@ -932,7 +932,6 @@ export async function createDataSourceWithoutProvider(
 
   const dataSourceView =
     await DataSourceViewResource.createDataSourceAndDefaultView(
-      auth,
       {
         name,
         description,
@@ -942,7 +941,8 @@ export async function createDataSourceWithoutProvider(
         assistantDefaultSelected: false,
         conversationId: conversation?.id,
       },
-      space
+      space,
+      auth.user()
     );
 
   try {

front/lib/resources/data_source_resource.ts

@@ -4,6 +4,7 @@ import type {
   DataSourceType,
   ModelId,
   Result,
+  UserType,
 } from "@dust-tt/types";
 import { formatUserFullName, Ok, removeNulls } from "@dust-tt/types";
 import type {
@@ -79,18 +80,18 @@ export class DataSourceResource extends ResourceWithSpace<DataSourceModel> {
   }
 
   static async makeNew(
-    auth: Authenticator,
     blob: Omit<
       CreationAttributes<DataSourceModel>,
       "editedAt" | "editedByUserId" | "vaultId"
     >,
     space: SpaceResource,
+    editedByUser?: UserType | null,
     transaction?: Transaction
   ) {
     const dataSource = await DataSourceModel.create(
       {
         ...blob,
-        editedByUserId: auth.user()?.id ?? null,
+        editedByUserId: editedByUser?.id ?? null,
         editedAt: new Date(),
         vaultId: space.id,
       },

front/lib/resources/data_source_view_resource.ts

@@ -7,6 +7,7 @@ import type {
   DataSourceViewType,
   ModelId,
   Result,
+  UserType,
 } from "@dust-tt/types";
 import { formatUserFullName, Ok, removeNulls } from "@dust-tt/types";
 import assert from "assert";
@@ -88,19 +89,19 @@ export class DataSourceViewResource extends ResourceWithSpace<DataSourceViewMode
   // Creation.
 
   private static async makeNew(
-    auth: Authenticator,
     blob: Omit<
       CreationAttributes<DataSourceViewModel>,
       "editedAt" | "editedByUserId" | "vaultId"
     >,
     space: SpaceResource,
     dataSource: DataSourceResource,
+    editedByUser?: UserType | null,
     transaction?: Transaction
   ) {
     const dataSourceView = await DataSourceViewResource.model.create(
       {
         ...blob,
-        editedByUserId: auth.user()?.id ?? null,
+        editedByUserId: editedByUser?.id ?? null,
         editedAt: new Date(),
         vaultId: space.id,
       },
@@ -117,54 +118,60 @@ export class DataSourceViewResource extends ResourceWithSpace<DataSourceViewMode
   }
 
   static async createDataSourceAndDefaultView(
-    auth: Authenticator,
     blob: Omit<CreationAttributes<DataSourceModel>, "editedAt" | "vaultId">,
-    space: SpaceResource
+    space: SpaceResource,
+    editedByUser?: UserType | null,
+    transaction?: Transaction
   ) {
-    return frontSequelize.transaction(async (transaction) => {
+    const createDataSourceAndView = async (t: Transaction) => {
       const dataSource = await DataSourceResource.makeNew(
-        auth,
         blob,
         space,
-        transaction
+        editedByUser,
+        t
       );
       return this.createDefaultViewInSpaceFromDataSourceIncludingAllDocuments(
-        auth,
-        dataSource.space,
+        space,
         dataSource,
-        transaction
+        editedByUser,
+        t
       );
-    });
+    };
+
+    if (transaction) {
+      return createDataSourceAndView(transaction);
+    }
+
+    return frontSequelize.transaction(createDataSourceAndView);
   }
 
   static async createViewInSpaceFromDataSource(
-    auth: Authenticator,
     space: SpaceResource,
     dataSource: DataSourceResource,
-    parentsIn: string[]
+    parentsIn: string[],
+    editedByUser?: UserType | null
   ) {
     return this.makeNew(
-      auth,
       {
         dataSourceId: dataSource.id,
         parentsIn,
         workspaceId: space.workspaceId,
         kind: "custom",
       },
       space,
-      dataSource
+      dataSource,
+      editedByUser
     );
   }
 
   // This view has access to all documents, which is represented by null.
   private static async createDefaultViewInSpaceFromDataSourceIncludingAllDocuments(
-    auth: Authenticator,
     space: SpaceResource,
     dataSource: DataSourceResource,
+    editedByUser?: UserType | null,
     transaction?: Transaction
   ) {
     return this.makeNew(
-      auth,
       {
         dataSourceId: dataSource.id,
         parentsIn: null,
@@ -173,6 +180,7 @@ export class DataSourceViewResource extends ResourceWithSpace<DataSourceViewMode
       },
       space,
       dataSource,
+      editedByUser,
       transaction
     );
   }

front/lib/resources/space_resource.ts

@@ -58,10 +58,11 @@ export class SpaceResource extends BaseResource<SpaceModel> {
 
   static async makeNew(
     blob: CreationAttributes<SpaceModel>,
-    groups: GroupResource[]
+    groups: GroupResource[],
+    transaction?: Transaction
   ) {
-    return frontSequelize.transaction(async (transaction) => {
-      const space = await SpaceModel.create(blob, { transaction });
+    const createSpace = async (t: Transaction) => {
+      const space = await SpaceModel.create(blob, { transaction: t });
 
       for (const group of groups) {
         await GroupSpaceModel.create(
@@ -70,12 +71,18 @@ export class SpaceResource extends BaseResource<SpaceModel> {
             vaultId: space.id,
             workspaceId: space.workspaceId,
           },
-          { transaction }
+          { transaction: t }
         );
       }
 
       return new this(SpaceModel, space.get(), groups);
-    });
+    };
+
+    if (transaction) {
+      return createSpace(transaction);
+    }
+
+    return frontSequelize.transaction(createSpace);
   }
 
   static async makeDefaultsForWorkspace(

front/migrations/20240730_backfill_data_source_views.ts

@@ -47,10 +47,10 @@ async function backfillDataSourceViewsForWorkspace(
 
     // Create a view for this data source in the global vault.
     await DataSourceViewResource.createViewInSpaceFromDataSource(
-      auth,
       globalVault,
       dataSource,
-      []
+      [],
+      auth.user()
     );
 
     updated++;

front/migrations/20240820_backfill_data_source_views.ts

@@ -39,10 +39,10 @@ async function backfillDefaultViewForDataSource(
 
   // Create a default view for this data source in the vault.
   await DataSourceViewResource.createViewInSpaceFromDataSource(
-    auth,
     vault,
     dataSource,
-    []
+    [],
+    auth.user()
   );
 
   logger.info(`View created for data source ${dataSource.id}.`);

front/migrations/20240821_backfill_all_data_source_views.ts

@@ -35,10 +35,10 @@ async function backfillDefaultViewForDataSource(
 
   // Create a default view for this data source in the vault.
   await DataSourceViewResource.createViewInSpaceFromDataSource(
-    auth,
     space,
     dataSource,
-    []
+    [],
+    auth.user()
   );
 
   logger.info(`View created for data source ${dataSource.id}.`);

front/pages/api/templates/[tId]/index.test.ts

@@ -2,8 +2,7 @@ import type { NextApiRequest, NextApiResponse } from "next";
 import { createMocks } from "node-mocks-http";
 import { describe, expect, vi } from "vitest";
 
-import { TemplateResource } from "@app/lib/resources/template_resource";
-import { templateFactory } from "@app/tests/utils/TemplateFactory";
+import { TemplateFactory } from "@app/tests/utils/TemplateFactory";
 import { itInTransaction } from "@app/tests/utils/utils";
 
 import handler from "./index";
@@ -55,11 +54,11 @@ describe("GET /api/templates/[tId]", () => {
   });
 
   itInTransaction("returns 404 when template is not published", async () => {
-    const template = await templateFactory().draft().create();
+    const template = await TemplateFactory.draft();
 
     const { req, res } = createMocks<NextApiRequest, NextApiResponse>({
       method: "GET",
-      query: { tId: TemplateResource.modelIdToSId(template) },
+      query: { tId: template.sId },
       headers: {},
     });
 
@@ -77,11 +76,11 @@ describe("GET /api/templates/[tId]", () => {
   itInTransaction(
     "returns template when it exists and is published",
     async () => {
-      const template = await templateFactory().published().create();
+      const template = await TemplateFactory.published();
 
       const { req, res } = createMocks<NextApiRequest, NextApiResponse>({
         method: "GET",
-        query: { tId: TemplateResource.modelIdToSId(template) },
+        query: { tId: template.sId },
         headers: {},
       });
 

front/pages/api/templates/index.test.ts

@@ -2,7 +2,7 @@ import type { NextApiRequest, NextApiResponse } from "next";
 import { createMocks } from "node-mocks-http";
 import { describe, expect, vi } from "vitest";
 
-import { templateFactory } from "@app/tests/utils/TemplateFactory";
+import { TemplateFactory } from "@app/tests/utils/TemplateFactory";
 import { itInTransaction } from "@app/tests/utils/utils";
 
 import handler from "./index";
@@ -40,9 +40,9 @@ describe("GET /api/templates", () => {
     });
 
     // Create test templates
-    const publishedTemplate1 = await templateFactory().published().create();
-    const publishedTemplate2 = await templateFactory().published().create();
-    await templateFactory().draft().create(); // Draft template should not be returned
+    const publishedTemplate1 = await TemplateFactory.published();
+    const publishedTemplate2 = await TemplateFactory.published();
+    await TemplateFactory.draft(); // Draft template should not be returned
 
     await handler(req, res);
 

front/pages/api/user/index.test.ts

@@ -1,6 +1,6 @@
 import { describe, expect } from "vitest";
 
-import { UserModel } from "@app/lib/resources/storage/models/user";
+import { UserResource } from "@app/lib/resources/user_resource";
 import { createPrivateApiMockRequest } from "@app/tests/utils/generic_private_api_tests";
 import { itInTransaction } from "@app/tests/utils/utils";
 
@@ -60,7 +60,7 @@ describe("GET /api/user", () => {
       success: true,
     });
 
-    const userAfterUpdate = await UserModel.findByPk(user.id);
+    const userAfterUpdate = await UserResource.fetchById(user.sId);
     expect(userAfterUpdate?.firstName).toBe("John");
     expect(userAfterUpdate?.lastName).toBe("Doe");
   });

front/pages/api/v1/w/[wId]/data_sources/index.test.ts

@@ -1,9 +1,9 @@
 import { describe, expect } from "vitest";
 
-import { dataSourceViewFactory } from "@app/tests/utils/DataSourceViewFactory";
+import { DataSourceViewFactory } from "@app/tests/utils/DataSourceViewFactory";
 import { createPublicApiMockRequest } from "@app/tests/utils/generic_public_api_tests";
-import { groupSpaceFactory } from "@app/tests/utils/GroupSpaceFactory";
-import { spaceFactory } from "@app/tests/utils/SpaceFactory";
+import { GroupSpaceFactory } from "@app/tests/utils/GroupSpaceFactory";
+import { SpaceFactory } from "@app/tests/utils/SpaceFactory";
 import {
   expectArrayOfObjectsWithSpecificLength,
   itInTransaction,
@@ -20,22 +20,22 @@ describe("GET /api/v1/w/[wId]/data_sources (legacy endpoint)", () => {
     expect(res._getStatusCode()).toBe(500);
   });
 
-  itInTransaction("returns data sources for the global space", async () => {
+  itInTransaction("returns data sources for the global space", async (t) => {
     const { req, res, workspace, globalGroup } =
       await createPublicApiMockRequest();
 
-    const space = await spaceFactory().global(workspace).create();
-    await groupSpaceFactory().associate(space, globalGroup);
+    const space = await SpaceFactory.global(workspace, t);
+    await GroupSpaceFactory.associate(space, globalGroup);
 
     // Create test data source views to the space
-    await dataSourceViewFactory().folder(workspace, space).create();
-    await dataSourceViewFactory().folder(workspace, space).create();
+    await DataSourceViewFactory.folder(workspace, space, t);
+    await DataSourceViewFactory.folder(workspace, space, t);
 
     // Create another space
-    const space2 = await spaceFactory().regular(workspace).create();
+    const space2 = await SpaceFactory.regular(workspace, t);
 
     // Create test data source views to the space (they should not be returned)
-    await dataSourceViewFactory().folder(workspace, space2).create();
+    await DataSourceViewFactory.folder(workspace, space2, t);
 
     // Execute request
     await handler(req, res);

front/pages/api/v1/w/[wId]/feature_flags.test.ts

@@ -1,14 +1,14 @@
 import { describe, expect, vi } from "vitest";
 
 import handler from "@app/pages/api/v1/w/[wId]/feature_flags";
-import { featureFlagFactory } from "@app/tests/utils/FeatureFlagFactory";
+import { FeatureFlagFactory } from "@app/tests/utils/FeatureFlagFactory";
 import {
   createPublicApiAuthenticationTests,
   createPublicApiMockRequest,
   createPublicApiSystemOnlyAuthenticationTests,
 } from "@app/tests/utils/generic_public_api_tests";
 import { itInTransaction } from "@app/tests/utils/utils";
-import { workspaceFactory } from "@app/tests/utils/WorkspaceFactory";
+import { WorkspaceFactory } from "@app/tests/utils/WorkspaceFactory";
 
 // Mock the getSession function to return the user without going through the auth0 session
 // Not sure to understand why it's not working with the generic_public_api_tests.ts mock
@@ -37,8 +37,8 @@ describe("GET /api/v1/w/[wId]/feature_flags", () => {
     });
 
     // Add features flag
-    await featureFlagFactory().basic("deepseek_feature", workspace).create();
-    await featureFlagFactory().basic("document_tracker", workspace).create();
+    await FeatureFlagFactory.basic("deepseek_feature", workspace);
+    await FeatureFlagFactory.basic("document_tracker", workspace);
 
     await handler(req, res);
 
@@ -97,10 +97,10 @@ describe("GET /api/v1/w/[wId]/feature_flags", () => {
         workspace: workspace1,
       } = await createPublicApiMockRequest({ systemKey: true });
 
-      const workspace2 = await workspaceFactory().basic().create();
+      const workspace2 = await WorkspaceFactory.basic();
 
-      await featureFlagFactory().basic("labs_trackers", workspace1).create();
-      await featureFlagFactory().basic("labs_transcripts", workspace2).create();
+      await FeatureFlagFactory.basic("labs_trackers", workspace1);
+      await FeatureFlagFactory.basic("labs_transcripts", workspace2);
 
       await handler(req, res);