twenty/packages/twenty-server/src/engine/twenty-orm/entity-manager/workspace-entity-manager.spec.ts

import {
  type FieldMetadataType,
  type ObjectsPermissions,
} from 'twenty-shared/types';
import { EntityManager } from 'typeorm';
import { EntityPersistExecutor } from 'typeorm/persistence/EntityPersistExecutor';
import { PlainObjectToDatabaseEntityTransformer } from 'typeorm/query-builder/transformer/PlainObjectToDatabaseEntityTransformer';

import { type WorkspaceInternalContext } from 'src/engine/twenty-orm/interfaces/workspace-internal-context.interface';

import { type WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { type FlatEntityMaps } from 'src/engine/metadata-modules/flat-entity/types/flat-entity-maps.type';
import { type FlatFieldMetadata } from 'src/engine/metadata-modules/flat-field-metadata/types/flat-field-metadata.type';
import { type FlatObjectMetadata } from 'src/engine/metadata-modules/flat-object-metadata/types/flat-object-metadata.type';
import { type GlobalWorkspaceDataSource } from 'src/engine/twenty-orm/global-workspace-datasource/global-workspace-datasource';
import { validateOperationIsPermittedOrThrow } from 'src/engine/twenty-orm/repository/permissions.utils';
import {
  setWorkspaceContext,
  withWorkspaceContext,
  type ORMWorkspaceContext,
} from 'src/engine/twenty-orm/storage/orm-workspace-context.storage';
import { getObjectMetadataFromEntityTarget } from 'src/engine/twenty-orm/utils/get-object-metadata-from-entity-target.util';

import { WorkspaceEntityManager } from './workspace-entity-manager';

jest.mock('src/engine/twenty-orm/repository/permissions.utils', () => ({
  validateOperationIsPermittedOrThrow: jest.fn(),
}));

jest.mock(
  'src/engine/twenty-orm/utils/get-object-metadata-from-entity-target.util',
  () => ({
    getObjectMetadataFromEntityTarget: jest.fn(),
  }),
);

jest.mock('src/engine/twenty-orm/utils/format-data.util', () => ({
  formatData: jest.fn().mockReturnValue([]),
}));

jest.mock('src/engine/twenty-orm/utils/format-result.util', () => ({
  formatResult: jest.fn().mockReturnValue([]),
}));

jest.mock(
  'src/engine/twenty-orm/entity-manager/workspace-entity-manager',
  () => ({
    ...jest.requireActual(
      'src/engine/twenty-orm/entity-manager/workspace-entity-manager',
    ),
  }),
);

const mockedWorkspaceUpdateQueryBuilder = {
  set: jest.fn().mockImplementation(() => ({
    where: jest.fn().mockReturnThis(),
    whereInIds: jest.fn().mockReturnThis(),
    execute: jest
      .fn()
      .mockResolvedValue({ affected: 1, raw: [], generatedMaps: [] }),
    returning: jest.fn().mockReturnThis(),
  })),
  execute: jest
    .fn()
    .mockResolvedValue({ affected: 1, raw: [], generatedMaps: [] }),
};

jest.mock('../repository/workspace-select-query-builder', () => ({
  WorkspaceSelectQueryBuilder: jest.fn().mockImplementation(() => ({
    where: jest.fn().mockReturnThis(),
    getMany: jest.fn().mockResolvedValue([]),
    getOne: jest.fn().mockResolvedValue(null),
    getManyAndCount: jest.fn().mockResolvedValue([[], 0]),
    execute: jest
      .fn()
      .mockResolvedValue({ affected: 1, raw: [], generatedMaps: [] }),
    setFindOptions: jest.fn().mockReturnThis(),
    returning: jest.fn().mockReturnThis(),
    update: jest.fn().mockReturnValue(mockedWorkspaceUpdateQueryBuilder),
    insert: jest.fn().mockReturnThis(),
  })),
}));

describe('WorkspaceEntityManager', () => {
  let entityManager: WorkspaceEntityManager;
  let mockDataSource: GlobalWorkspaceDataSource;
  let mockPermissionOptions: {
    shouldBypassPermissionChecks: boolean;
    objectRecordsPermissions?: ObjectsPermissions;
  };
  let mockInternalContext: WorkspaceInternalContext;
  let mockWorkspaceContext: ORMWorkspaceContext;

  beforeEach(() => {
    const mockFlatObjectMetadata: FlatObjectMetadata = {
      id: 'test-entity-id',
      nameSingular: 'test-entity',
      namePlural: 'test-entities',
      labelSingular: 'Test Entity',
      labelPlural: 'Test Entities',
      workspaceId: 'test-workspace-id',
      icon: 'test-icon',
      color: null,
      isCustom: false,
      isRemote: false,
      isAuditLogged: false,
      isSearchable: false,
      isSystem: false,
      isActive: true,
      targetTableName: 'test_entity',
      fieldIds: ['field-id'],
      indexMetadataIds: [],
      objectPermissionIds: [],
      fieldPermissionIds: [],
      viewIds: [],
      universalIdentifier: 'test-entity-id',
      description: null,
      imageIdentifierFieldMetadataId: null,
      labelIdentifierFieldMetadataId: null,
      shortcut: null,
      standardOverrides: null,
      applicationId: 'test-application-id',
      isLabelSyncedWithName: false,
      isUIReadOnly: false,
      duplicateCriteria: null,
      createdAt: new Date().toISOString(),
      updatedAt: new Date().toISOString(),
      applicationUniversalIdentifier: 'test-application-id',
      fieldUniversalIdentifiers: ['field-id'],
      objectPermissionUniversalIdentifiers: [],
      fieldPermissionUniversalIdentifiers: [],
      viewUniversalIdentifiers: [],
      indexMetadataUniversalIdentifiers: [],
      labelIdentifierFieldMetadataUniversalIdentifier: null,
      imageIdentifierFieldMetadataUniversalIdentifier: null,
    };

    (getObjectMetadataFromEntityTarget as jest.Mock).mockReturnValue(
      mockFlatObjectMetadata,
    );

    const mockFlatFieldMetadata: FlatFieldMetadata = {
      id: 'field-id',
      type: 'TEXT' as FieldMetadataType,
      name: 'fieldName',
      label: 'Field Name',
      objectMetadataId: 'test-entity-id',
      isNullable: true,
      isLabelSyncedWithName: false,
      createdAt: new Date().toISOString(),
      updatedAt: new Date().toISOString(),
      universalIdentifier: 'field-id',
      defaultValue: null,
      description: null,
      icon: null,
      isActive: true,
      isCustom: false,
      isSystem: false,
      isUIReadOnly: false,
      isUnique: false,
      options: null,
      settings: null,
      standardOverrides: null,
      workspaceId: 'test-workspace-id',
      viewFieldIds: [],
      viewFilterIds: [],
      fieldPermissionIds: [],
      kanbanAggregateOperationViewIds: [],
      calendarViewIds: [],
      mainGroupByFieldMetadataViewIds: [],
      relationTargetFieldMetadataId: null,
      relationTargetObjectMetadataId: null,
      morphId: null,
      applicationId: 'application-id',
      applicationUniversalIdentifier: 'application-id',
      objectMetadataUniversalIdentifier: 'test-entity-id',
      relationTargetObjectMetadataUniversalIdentifier: null,
      relationTargetFieldMetadataUniversalIdentifier: null,
      viewFilterUniversalIdentifiers: [],
      viewFieldUniversalIdentifiers: [],
      kanbanAggregateOperationViewUniversalIdentifiers: [],
      calendarViewUniversalIdentifiers: [],
      mainGroupByFieldMetadataViewUniversalIdentifiers: [],
      fieldPermissionUniversalIdentifiers: [],
      viewSortIds: [],
      viewSortUniversalIdentifiers: [],
      universalSettings: null,
    };

    const flatObjectMetadataMaps: FlatEntityMaps<FlatObjectMetadata> = {
      byUniversalIdentifier: {
        'test-entity-id': mockFlatObjectMetadata,
      },
      universalIdentifierById: {
        'test-entity-id': 'test-entity-id',
      },
      universalIdentifiersByApplicationId: {},
    };

    const flatFieldMetadataMaps: FlatEntityMaps<FlatFieldMetadata> = {
      byUniversalIdentifier: {
        'field-id': mockFlatFieldMetadata,
      },
      universalIdentifierById: {
        'field-id': 'field-id',
      },
      universalIdentifiersByApplicationId: {},
    };

    mockInternalContext = {
      workspaceId: 'test-workspace-id',
      flatObjectMetadataMaps,
      flatFieldMetadataMaps,
      flatIndexMaps: {
        byUniversalIdentifier: {},
        universalIdentifierById: {},
        universalIdentifiersByApplicationId: {},
      },
      flatRowLevelPermissionPredicateMaps: {
        byUniversalIdentifier: {},
        universalIdentifierById: {},
        universalIdentifiersByApplicationId: {},
      },
      flatRowLevelPermissionPredicateGroupMaps: {
        byUniversalIdentifier: {},
        universalIdentifierById: {},
        universalIdentifiersByApplicationId: {},
      },
      objectIdByNameSingular: {
        'test-entity': 'test-entity-id',
      },
      featureFlagsMap: {
        IS_UNIQUE_INDEXES_ENABLED: false,
        IS_JSON_FILTER_ENABLED: false,
        IS_MARKETPLACE_SETTING_TAB_VISIBLE: false,
        IS_PUBLIC_DOMAIN_ENABLED: false,
        IS_EMAIL_GROUP_ENABLED: false,
        IS_JUNCTION_RELATIONS_ENABLED: false,
        IS_REST_METADATA_API_NEW_FORMAT_DIRECT: false,
        IS_LOGIC_FUNCTION_PREBUILT_MODE_ENABLED: false,
        IS_SETTINGS_DISCOVERY_HERO_ENABLED: false,
        IS_WORKFLOW_RUN_STEP_LOGS_ENABLED: false,
      },
      userWorkspaceRoleMap: {},
      eventEmitterService: {
        emitMutationEvent: jest.fn(),
        emitDatabaseBatchEvent: jest.fn(),
        emitCustomBatchEvent: jest.fn(),
      } as any,
      coreDataSource: {
        getRepository: jest.fn(() => ({
          find: jest.fn(),
          softDelete: jest.fn(),
        })),
      } as any,
    } as WorkspaceInternalContext;

    mockDataSource = {
      featureFlagMap: {
        IS_UNIQUE_INDEXES_ENABLED: false,
        IS_JSON_FILTER_ENABLED: false,
        IS_PUBLIC_DOMAIN_ENABLED: false,
      },
      permissionsPerRoleId: {},
      eventEmitterService: mockInternalContext.eventEmitterService,
      coreDataSource: mockInternalContext.coreDataSource,
    } as GlobalWorkspaceDataSource;

    mockPermissionOptions = {
      shouldBypassPermissionChecks: false,
      objectRecordsPermissions: {
        'test-entity': {
          canReadObjectRecords: true,
          canUpdateObjectRecords: false,
          canSoftDeleteObjectRecords: false,
          canDestroyObjectRecords: false,
          restrictedFields: {},
          rowLevelPermissionPredicates: [],
          rowLevelPermissionPredicateGroups: [],
        },
      },
    };

    const mockAuthContext = {
      user: { id: 'user-id' },
      workspace: { id: 'test-workspace-id' },
      workspaceMemberId: 'workspace-member-id',
      userWorkspaceId: 'user-workspace-id',
      apiKey: null,
    } as unknown as WorkspaceAuthContext;

    mockWorkspaceContext = {
      authContext: mockAuthContext,
      flatObjectMetadataMaps,
      flatFieldMetadataMaps,
      flatIndexMaps: mockInternalContext.flatIndexMaps,
      flatRowLevelPermissionPredicateMaps:
        mockInternalContext.flatRowLevelPermissionPredicateMaps,
      flatRowLevelPermissionPredicateGroupMaps:
        mockInternalContext.flatRowLevelPermissionPredicateGroupMaps,
      objectIdByNameSingular: mockInternalContext.objectIdByNameSingular,
      featureFlagsMap: mockInternalContext.featureFlagsMap,
      permissionsPerRoleId: mockDataSource.permissionsPerRoleId,
      entityMetadatas: [],
      userWorkspaceRoleMap: {
        'user-workspace-id': 'role-id',
      },
      apiKeyRoleMap: {},
    };

    setWorkspaceContext(mockWorkspaceContext);

    // Mock TypeORM connection methods
    const mockWorkspaceDataSource = {
      getMetadata: jest.fn().mockReturnValue({
        name: 'test-entity',
        columns: [],
        relations: [],
        findInheritanceMetadata: jest.fn(),
        findColumnWithPropertyPath: jest.fn(),
      }),
      eventEmitterService: mockInternalContext.eventEmitterService,
      coreDataSource: mockInternalContext.coreDataSource,
      createQueryBuilder: jest.fn().mockReturnValue({
        delete: jest.fn().mockReturnThis(),
        from: jest.fn().mockReturnThis(),
        where: jest.fn().mockReturnThis(),
        execute: jest
          .fn()
          .mockResolvedValue({ affected: 1, raw: [], generatedMaps: [] }),
        getMany: jest.fn().mockResolvedValue([]),
        getOne: jest.fn().mockResolvedValue(null),
        getManyAndCount: jest.fn().mockResolvedValue([[], 0]),
        update: jest.fn().mockReturnThis(),
        softDelete: jest.fn().mockReturnThis(),
        restore: jest.fn().mockReturnThis(),
      }),
      createQueryRunner: jest.fn().mockReturnValue({
        connect: jest.fn(),
        startTransaction: jest.fn(),
        commitTransaction: jest.fn(),
        rollbackTransaction: jest.fn(),
        release: jest.fn(),
        clearTable: jest.fn(),
      }),
      createQueryRunnerForEntityPersistExecutor: jest.fn().mockReturnValue({
        connect: jest.fn(),
        startTransaction: jest.fn(),
        commitTransaction: jest.fn(),
        rollbackTransaction: jest.fn(),
        release: jest.fn(),
        clearTable: jest.fn(),
      }),
    };

    entityManager = new WorkspaceEntityManager(mockDataSource);

    Object.defineProperty(entityManager, 'connection', {
      get: () => mockWorkspaceDataSource,
    });

    jest.spyOn(entityManager as any, 'validatePermissions');
    jest.spyOn(entityManager as any, 'createQueryBuilder');
    jest
      .spyOn(entityManager as any, 'getFormattedResultWithoutNonReadableFields')
      .mockImplementation(
        ({ formattedResult }: { formattedResult: string[] }) => formattedResult,
      );

    jest
      .spyOn(entityManager as any, 'extractTargetNameSingularFromEntityTarget')
      .mockImplementation((entityName: string) => {
        return entityName;
      });

    // Mock typeORM's EntityManager methods
    jest
      .spyOn(EntityManager.prototype, 'save')
      .mockImplementation(() => Promise.resolve({}));
    jest
      .spyOn(EntityManager.prototype, 'update')
      .mockImplementation(() =>
        Promise.resolve({ affected: 1, raw: [], generatedMaps: [] }),
      );
    jest
      .spyOn(EntityManager.prototype, 'restore')
      .mockImplementation(() =>
        Promise.resolve({ affected: 1, raw: [], generatedMaps: [] }),
      );
    jest
      .spyOn(EntityManager.prototype, 'clear')
      .mockImplementation(() => Promise.resolve());
    jest
      .spyOn(EntityManager.prototype, 'preload')
      .mockImplementation(() => Promise.resolve({}));

    jest
      .spyOn(EntityPersistExecutor.prototype, 'execute')
      .mockImplementation(() => Promise.resolve());

    jest
      .spyOn(PlainObjectToDatabaseEntityTransformer.prototype, 'transform')
      .mockImplementation(() => Promise.resolve({}));

    // Mock metadata methods
    const mockMetadata = {
      hasAllPrimaryKeys: jest.fn().mockReturnValue(true),
      columns: [],
      relations: [],
      findInheritanceMetadata: jest.fn(),
      findColumnWithPropertyPath: jest.fn(),
    };

    // Update mockWorkspaceDataSource to include metadata
    mockWorkspaceDataSource.getMetadata = jest
      .fn()
      .mockReturnValue(mockMetadata);

    // Reset the mock before each test
    jest.clearAllMocks();
  });

  describe('Query Method', () => {
    it('should call validatePermissions and validateOperationIsPermittedOrThrow for find', async () => {
      await withWorkspaceContext(mockWorkspaceContext, () =>
        entityManager.find('test-entity', {}, mockPermissionOptions),
      );

      expect(entityManager.createQueryBuilder).toHaveBeenCalledWith(
        'test-entity',
        undefined,
        undefined,
        mockPermissionOptions,
      );
    });
    it('should throw error for query', async () => {
      expect(() => entityManager.query('SELECT * FROM test')).toThrow(
        'Method not allowed.',
      );
    });
  });

  describe('Save Methods', () => {
    it('should call validatePermissions and validateOperationIsPermittedOrThrow for save', async () => {
      await withWorkspaceContext(mockWorkspaceContext, () =>
        entityManager.save(
          'test-entity',
          {},
          { reload: false },
          mockPermissionOptions,
        ),
      );
      expect(entityManager['validatePermissions']).toHaveBeenCalledWith({
        target: 'test-entity',
        operationType: 'update',
        permissionOptions: mockPermissionOptions,
        selectedColumns: [],
        updatedColumns: [],
      });
      expect(validateOperationIsPermittedOrThrow).toHaveBeenCalledWith({
        entityName: 'test-entity',
        operationType: 'update',
        flatObjectMetadataMaps: mockInternalContext.flatObjectMetadataMaps,
        flatFieldMetadataMaps: mockInternalContext.flatFieldMetadataMaps,
        objectIdByNameSingular: mockInternalContext.objectIdByNameSingular,
        objectsPermissions: mockPermissionOptions.objectRecordsPermissions,
        selectedColumns: [],
        allFieldsSelected: false,
        updatedColumns: [],
      });
    });
  });

  describe('Update Methods', () => {
    it('should call createQueryBuilder with permissionOptions for update', async () => {
      await withWorkspaceContext(mockWorkspaceContext, () =>
        entityManager.update('test-entity', {}, {}, mockPermissionOptions),
      );
      expect(entityManager['createQueryBuilder']).toHaveBeenCalledWith(
        'test-entity',
        undefined,
        undefined,
        mockPermissionOptions,
      );
    });
  });

  describe('Other Methods', () => {
    it('should call validatePermissions and validateOperationIsPermittedOrThrow for clear', async () => {
      await withWorkspaceContext(mockWorkspaceContext, () =>
        entityManager.clear('test-entity', mockPermissionOptions),
      );
      expect(entityManager['validatePermissions']).toHaveBeenCalledWith({
        target: 'test-entity',
        operationType: 'delete',
        permissionOptions: mockPermissionOptions,
        selectedColumns: [],
      });
      expect(validateOperationIsPermittedOrThrow).toHaveBeenCalledWith({
        entityName: 'test-entity',
        operationType: 'delete',
        flatObjectMetadataMaps: mockInternalContext.flatObjectMetadataMaps,
        flatFieldMetadataMaps: mockInternalContext.flatFieldMetadataMaps,
        objectIdByNameSingular: mockInternalContext.objectIdByNameSingular,
        objectsPermissions: mockPermissionOptions.objectRecordsPermissions,
        selectedColumns: [],
        allFieldsSelected: false,
        updatedColumns: [],
      });
    });
  });
});