upstream-twenty/packages/twenty-server/test/integration/upgrade/utils/upgrade-sequence-runner-integration-test.util.ts

import { Test, type TestingModule } from '@nestjs/testing';
import { getDataSourceToken, getRepositoryToken } from '@nestjs/typeorm';

import { config } from 'dotenv';
import { WorkspaceActivationStatus } from 'twenty-shared/workspace';
import { DataSource, type Repository } from 'typeorm';

import { CommandShutdownService } from 'src/database/commands/command-runners/command-shutdown.service';
import { WorkspaceIteratorService } from 'src/database/commands/command-runners/workspace-iterator.service';
import { CoreEntityCacheService } from 'src/engine/core-entity-cache/services/core-entity-cache.service';
import { ApplicationEntity } from 'src/engine/core-modules/application/application.entity';
import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';
import { InstanceCommandRunnerService } from 'src/engine/core-modules/upgrade/services/instance-command-runner.service';
import { UpgradeMigrationService } from 'src/engine/core-modules/upgrade/services/upgrade-migration.service';
import {
  UpgradeSequenceReaderService,
  type UpgradeStep,
  type WorkspaceUpgradeStep,
} from 'src/engine/core-modules/upgrade/services/upgrade-sequence-reader.service';
import { UpgradeSequenceRunnerService } from 'src/engine/core-modules/upgrade/services/upgrade-sequence-runner.service';
import { UpgradeStatusCacheService } from 'src/engine/core-modules/upgrade/services/upgrade-status-cache.service';
import { UpgradeStatusService } from 'src/engine/core-modules/upgrade/services/upgrade-status.service';
import { WorkspaceCommandRunnerService } from 'src/engine/core-modules/upgrade/services/workspace-command-runner.service';
import { UpgradeMigrationEntity } from 'src/engine/core-modules/upgrade/upgrade-migration.entity';
import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { UpgradeAwareEntityMetadataAdapter } from 'src/engine/twenty-orm/upgrade-aware/upgrade-aware-entity-metadata.adapter';
import {
  SEED_APPLE_WORKSPACE_ID,
  SEED_EMPTY_WORKSPACE_3_ID,
  SEED_EMPTY_WORKSPACE_4_ID,
  SEED_YCOMBINATOR_WORKSPACE_ID,
} from 'src/engine/workspace-manager/dev-seeder/core/constants/seeder-workspaces.constant';
import { createWorkspace } from 'src/engine/workspace-manager/dev-seeder/core/utils/seed-workspace.util';
import { WorkspaceVersionService } from 'src/engine/workspace-manager/workspace-version/services/workspace-version.service';

jest.useRealTimers();

config({
  path: process.env.NODE_ENV === 'test' ? '.env.test' : '.env',
  override: true,
});

export const WS_1 = SEED_APPLE_WORKSPACE_ID;
export const WS_2 = SEED_YCOMBINATOR_WORKSPACE_ID;
export const WS_3 = SEED_EMPTY_WORKSPACE_3_ID;
export const WS_4 = SEED_EMPTY_WORKSPACE_4_ID;

const FK_WORKSPACE_FIXTURES = [
  {
    workspaceId: WS_3,
    applicationId: 'f1c0ffee-0000-4000-8000-0000000000c3',
    subdomain: 'upgrade-test-fixture-3',
  },
  {
    workspaceId: WS_4,
    applicationId: 'f1c0ffee-0000-4000-8000-0000000000c4',
    subdomain: 'upgrade-test-fixture-4',
  },
];

const seedEmptyWorkspaces = async (dataSource: DataSource) => {
  const queryRunner = dataSource.createQueryRunner();

  await queryRunner.connect();
  await queryRunner.startTransaction();

  try {
    for (const {
      workspaceId,
      applicationId,
      subdomain,
    } of FK_WORKSPACE_FIXTURES) {
      await createWorkspace({
        queryRunner,
        schemaName: 'core',
        createWorkspaceInput: {
          id: workspaceId,
          displayName: subdomain,
          subdomain,
          inviteHash: `${subdomain}.dev-invite-hash`,
          logo: '',
          activationStatus: WorkspaceActivationStatus.PENDING_CREATION,
          isTwoFactorAuthenticationEnforced: false,
          workspaceCustomApplicationId: applicationId,
        },
      });

      await queryRunner.manager
        .createQueryBuilder()
        .insert()
        .into(ApplicationEntity)
        .values({
          id: applicationId,
          universalIdentifier: applicationId,
          name: 'upgrade-test-fixture',
          sourcePath: '',
          workspaceId,
        })
        .orIgnore()
        .execute();
    }

    await queryRunner.commitTransaction();
  } catch (error) {
    await queryRunner.rollbackTransaction();
    throw error;
  } finally {
    await queryRunner.release();
  }
};

const EXECUTED_BY_VERSION = '42.42.42';

const noopAsync = async () => {};

const DEFAULT_STEP_VERSION = '1.21.0';

export const makeStep = (
  kind: UpgradeStep['kind'],
  name: string,
  version: string = DEFAULT_STEP_VERSION,
): UpgradeStep => {
  const command =
    kind === 'workspace'
      ? { runOnWorkspace: noopAsync }
      : kind === 'slow-instance'
        ? { up: noopAsync, down: noopAsync, runDataMigration: noopAsync }
        : { up: noopAsync, down: noopAsync };

  return {
    kind,
    name,
    command,
    version,
    timestamp: 0,
  } as unknown as UpgradeStep;
};

export const makeFastInstance = (name: string, version?: string) =>
  makeStep('fast-instance', name, version);

export const makeSlowInstance = (name: string, version?: string) =>
  makeStep('slow-instance', name, version);

export const makeWorkspace = (name: string, version?: string) =>
  makeStep('workspace', name, version) as WorkspaceUpgradeStep;

// Steps the status service can resolve a version from: it reads the version
// off the command name, not off the step's `version` field.
export const makeVersionedStep = (
  kind: UpgradeStep['kind'],
  { version, label }: { version: string; label: string },
): UpgradeStep => makeStep(kind, `${version}_${label}_0`, version);

let mockActiveWorkspaceIds: string[] = [];

export const setMockActiveWorkspaceIds = (ids: string[]) => {
  mockActiveWorkspaceIds = ids;
};

export const DEFAULT_OPTIONS = {
  workspaceIds: undefined,
  startFromWorkspaceId: undefined,
  workspaceCountLimit: undefined,
  dryRun: false,
  verbose: false,
};

type IntegrationTestModule = Awaited<
  ReturnType<typeof createUpgradeSequenceRunnerIntegrationTestModule>
>;

export type IntegrationTestContext = {
  [K in keyof IntegrationTestModule]: IntegrationTestModule[K];
};

export const createUpgradeSequenceRunnerIntegrationTestModule = async () => {
  const dataSource = new DataSource({
    type: 'postgres',
    url: process.env.PG_DATABASE_URL,
    schema: 'core',
    entities: [
      'src/engine/core-modules/**/*.entity.ts',
      'src/engine/metadata-modules/**/*.entity.ts',
    ],
    synchronize: false,
  });

  await dataSource.initialize();

  await seedEmptyWorkspaces(dataSource);

  const migrationRepo: Repository<UpgradeMigrationEntity> =
    dataSource.getRepository(UpgradeMigrationEntity);

  const module: TestingModule = await Test.createTestingModule({
    providers: [
      {
        provide: getRepositoryToken(UpgradeMigrationEntity),
        useValue: migrationRepo,
      },
      {
        provide: getDataSourceToken(),
        useValue: dataSource,
      },
      {
        provide: TwentyConfigService,
        useValue: {
          get: (key: string) =>
            key === 'APP_VERSION' ? EXECUTED_BY_VERSION : undefined,
        },
      },
      UpgradeMigrationService,
      {
        provide: WorkspaceVersionService,
        useValue: {
          getProvisionedWorkspaceIds: jest
            .fn()
            .mockImplementation(async () => mockActiveWorkspaceIds),
          hasProvisionedWorkspaces: jest
            .fn()
            .mockImplementation(async () => mockActiveWorkspaceIds.length > 0),
        },
      },
      {
        provide: UpgradeSequenceReaderService,
        useFactory: () => new UpgradeSequenceReaderService({} as any),
      },
      {
        provide: getRepositoryToken(WorkspaceEntity),
        useValue: dataSource.getRepository(WorkspaceEntity),
      },
      {
        provide: UpgradeStatusCacheService,
        useValue: {
          getComputedAt: jest.fn().mockResolvedValue(null),
          getBehindWorkspaceIds: jest.fn().mockResolvedValue([]),
          getFailedWorkspaceIds: jest.fn().mockResolvedValue([]),
          getUpToDateWorkspaceCount: jest.fn().mockResolvedValue(0),
          write: jest.fn().mockResolvedValue(undefined),
          invalidate: jest.fn().mockResolvedValue(undefined),
        },
      },
      {
        provide: CoreEntityCacheService,
        useValue: { get: jest.fn().mockResolvedValue(null) },
      },
      UpgradeStatusService,
      InstanceCommandRunnerService,
      WorkspaceCommandRunnerService,
      {
        provide: WorkspaceIteratorService,
        useValue: {
          iterate: jest.fn().mockImplementation(async (args: any) => {
            const { callback, workspaceIds } = args;
            const ids = workspaceIds ?? [WS_1];
            const report = {
              fail: [] as any[],
              success: [] as any[],
              interrupted: false,
            };

            for (const [index, workspaceId] of ids.entries()) {
              try {
                await callback({
                  workspaceId,
                  index,
                  total: ids.length,
                  dataSource,
                });
                report.success.push({ workspaceId });
              } catch (error) {
                report.fail.push({ error, workspaceId });
              }
            }

            return report;
          }),
        },
      },
      {
        provide: UpgradeAwareEntityMetadataAdapter,
        useValue: {
          refresh: jest.fn().mockResolvedValue(undefined),
          isEntityAvailable: jest.fn().mockReturnValue(true),
          getHiddenColumnPropertyNames: jest.fn().mockReturnValue(new Set()),
        },
      },
      CommandShutdownService,
      UpgradeSequenceRunnerService,
    ],
  }).compile();

  const runner = module.get(UpgradeSequenceRunnerService);

  jest.spyOn(runner['logger'], 'log').mockImplementation();
  jest.spyOn(runner['logger'], 'error').mockImplementation();
  jest.spyOn(runner['logger'], 'warn').mockImplementation();

  const instanceCommandRunnerService = module.get(InstanceCommandRunnerService);

  jest
    .spyOn(instanceCommandRunnerService['logger'], 'log')
    .mockImplementation();
  jest
    .spyOn(instanceCommandRunnerService['logger'], 'error')
    .mockImplementation();

  const workspaceCommandRunnerService = module.get(
    WorkspaceCommandRunnerService,
  );

  jest
    .spyOn(workspaceCommandRunnerService['logger'], 'log')
    .mockImplementation();
  jest
    .spyOn(workspaceCommandRunnerService['logger'], 'error')
    .mockImplementation();

  return {
    module,
    dataSource,
    runner,
    upgradeStatusService: module.get(UpgradeStatusService),
    upgradeSequenceReaderService: module.get(UpgradeSequenceReaderService),
  };
};

let seedSequenceCounter = 0;

export const resetSeedSequenceCounter = () => {
  seedSequenceCounter = 0;
};

export const seedInstanceMigration = async (
  dataSource: DataSource,
  {
    name,
    status,
    workspaceIds = [],
    attempt = 1,
  }: {
    name: string;
    status: 'completed' | 'failed';
    workspaceIds?: string[];
    attempt?: number;
  },
) => {
  // Seeds must have past timestamps so the runner's NOW()-based records
  // always sort after them in createdAt order.
  const createdAt = new Date(
    Date.now() - (1000000 - seedSequenceCounter * 1000),
  ).toISOString();

  seedSequenceCounter++;

  const values: string[] = [];
  const args: unknown[] = [];
  let paramIndex = 1;

  values.push(
    `($${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++}, NULL, $${paramIndex++}, false)`,
  );
  args.push(name, status, attempt, EXECUTED_BY_VERSION, createdAt);

  for (const workspaceId of workspaceIds) {
    values.push(
      `($${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++}, $${paramIndex++}, false)`,
    );
    args.push(
      name,
      status,
      attempt,
      EXECUTED_BY_VERSION,
      workspaceId,
      createdAt,
    );
  }

  await dataSource.query(
    `INSERT INTO core."upgradeMigration" (name, status, attempt, "executedByVersion", "workspaceId", "createdAt", "isInitial")
     VALUES ${values.join(', ')}`,
    args,
  );
};

export const seedWorkspaceMigration = async (
  dataSource: DataSource,
  {
    name,
    status,
    workspaceId,
    attempt = 1,
    isInitial = false,
    useCurrentTimestamp = false,
  }: {
    name: string;
    status: 'completed' | 'failed';
    workspaceId: string;
    attempt?: number;
    isInitial?: boolean;
    useCurrentTimestamp?: boolean;
  },
) => {
  if (useCurrentTimestamp) {
    await dataSource.query(
      `INSERT INTO core."upgradeMigration" (name, status, attempt, "executedByVersion", "workspaceId", "isInitial")
       VALUES ($1, $2, $3, $4, $5, $6)`,
      [name, status, attempt, EXECUTED_BY_VERSION, workspaceId, isInitial],
    );
  } else {
    const createdAt = new Date(
      Date.now() - (1000000 - seedSequenceCounter * 1000),
    ).toISOString();

    seedSequenceCounter++;

    await dataSource.query(
      `INSERT INTO core."upgradeMigration" (name, status, attempt, "executedByVersion", "workspaceId", "createdAt", "isInitial")
       VALUES ($1, $2, $3, $4, $5, $6, $7)`,
      [
        name,
        status,
        attempt,
        EXECUTED_BY_VERSION,
        workspaceId,
        createdAt,
        isInitial,
      ],
    );
  }
};

export const snapshotUpgradeMigrations = async (
  dataSource: DataSource,
): Promise<UpgradeMigrationEntity[]> =>
  dataSource.query(
    `SELECT id, name, status, attempt, "executedByVersion", "errorMessage", "isInitial", "workspaceId", "createdAt"
     FROM core."upgradeMigration"`,
  );

export const restoreUpgradeMigrations = async (
  dataSource: DataSource,
  rows: UpgradeMigrationEntity[],
): Promise<void> => {
  await dataSource.query('DELETE FROM core."upgradeMigration"');

  if (rows.length === 0) {
    return;
  }

  const columnsPerRow = 9;
  const valueGroups: string[] = [];
  const args: unknown[] = [];
  let paramIndex = 1;

  for (const row of rows) {
    const placeholders = Array.from(
      { length: columnsPerRow },
      () => `$${paramIndex++}`,
    );

    valueGroups.push(`(${placeholders.join(', ')})`);
    args.push(
      row.id,
      row.name,
      row.status,
      row.attempt,
      row.executedByVersion,
      row.errorMessage,
      row.isInitial,
      row.workspaceId,
      row.createdAt,
    );
  }

  await dataSource.query(
    `INSERT INTO core."upgradeMigration"
       (id, name, status, attempt, "executedByVersion", "errorMessage", "isInitial", "workspaceId", "createdAt")
     VALUES ${valueGroups.join(', ')}`,
    args,
  );
};

export type ExecutedMigrationRecord = {
  name: string;
  status: string;
  attempt: number;
  workspaceId: string | null;
  isInitial: boolean;
};

export const testGetExecutedMigrationsInOrder = async (
  dataSource: DataSource,
): Promise<ExecutedMigrationRecord[]> => {
  return dataSource.query(
    `SELECT name, status, attempt, "workspaceId", "isInitial"
     FROM core."upgradeMigration"
     ORDER BY "createdAt" ASC, "workspaceId" ASC NULLS FIRST, attempt ASC`,
  );
};

export const migrationRecordToKey = ({
  name,
  workspaceId,
  status,
  attempt,
  isInitial,
}: ExecutedMigrationRecord): string => {
  const scope = workspaceId ?? 'instance';
  const initial = isInitial ? ':initial' : '';

  return `${name}:${scope}:${status}:${attempt}${initial}`;
};