AFFiNE/packages/backend/server/tests/workspace-blobs.spec.ts

import type { INestApplication } from '@nestjs/common';
import { Test } from '@nestjs/testing';
import { PrismaClient } from '@prisma/client';
import test from 'ava';
import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs';
import request from 'supertest';

import { AppModule } from '../src/app';
import {
  collectMigrations,
  RevertCommand,
  RunCommand,
} from '../src/data/commands/run';
import { QuotaService, QuotaType } from '../src/modules/quota';
import {
  checkBlobSize,
  collectAllBlobSizes,
  createWorkspace,
  getWorkspaceBlobsSize,
  listBlobs,
  setBlob,
  signUp,
} from './utils';

let app: INestApplication;
let quota: QuotaService;

const client = new PrismaClient();

// cleanup database before each test
test.beforeEach(async () => {
  await client.$connect();
  await client.user.deleteMany({});
  await client.snapshot.deleteMany({});
  await client.update.deleteMany({});
  await client.workspace.deleteMany({});
  await client.$disconnect();
});

test.beforeEach(async () => {
  const module = await Test.createTestingModule({
    imports: [AppModule],
    providers: [RevertCommand, RunCommand],
  }).compile();
  app = module.createNestApplication();
  app.use(
    graphqlUploadExpress({
      maxFileSize: 10 * 1024 * 1024,
      maxFiles: 5,
    })
  );
  quota = module.get(QuotaService);

  // init features
  const run = module.get(RunCommand);
  const revert = module.get(RevertCommand);
  const migrations = await collectMigrations();
  await Promise.allSettled(migrations.map(m => revert.run([m.name])));
  await run.run();

  await app.init();
});

test.afterEach.always(async () => {
  await app.close();
});

test('should set blobs', async t => {
  const u1 = await signUp(app, 'u1', 'u1@affine.pro', '1');

  const workspace = await createWorkspace(app, u1.token.token);

  const buffer1 = Buffer.from([0, 0]);
  const hash1 = await setBlob(app, u1.token.token, workspace.id, buffer1);
  const buffer2 = Buffer.from([0, 1]);
  const hash2 = await setBlob(app, u1.token.token, workspace.id, buffer2);

  const server = app.getHttpServer();

  const response1 = await request(server)
    .get(`/api/workspaces/${workspace.id}/blobs/${hash1}`)
    .auth(u1.token.token, { type: 'bearer' })
    .buffer();

  t.deepEqual(response1.body, buffer1, 'failed to get blob');

  const response2 = await request(server)
    .get(`/api/workspaces/${workspace.id}/blobs/${hash2}`)
    .auth(u1.token.token, { type: 'bearer' })
    .buffer();

  t.deepEqual(response2.body, buffer2, 'failed to get blob');
});

test('should list blobs', async t => {
  const u1 = await signUp(app, 'u1', 'u1@affine.pro', '1');

  const workspace = await createWorkspace(app, u1.token.token);
  const blobs = await listBlobs(app, u1.token.token, workspace.id);
  t.is(blobs.length, 0, 'failed to list blobs');

  const buffer1 = Buffer.from([0, 0]);
  const hash1 = await setBlob(app, u1.token.token, workspace.id, buffer1);
  const buffer2 = Buffer.from([0, 1]);
  const hash2 = await setBlob(app, u1.token.token, workspace.id, buffer2);

  const ret = await listBlobs(app, u1.token.token, workspace.id);
  t.is(ret.length, 2, 'failed to list blobs');
  t.is(ret[0], hash1, 'failed to list blobs');
  t.is(ret[1], hash2, 'failed to list blobs');
});

test('should calc blobs size', async t => {
  const u1 = await signUp(app, 'u1', 'u1@affine.pro', '1');

  const workspace = await createWorkspace(app, u1.token.token);

  const buffer1 = Buffer.from([0, 0]);
  await setBlob(app, u1.token.token, workspace.id, buffer1);
  const buffer2 = Buffer.from([0, 1]);
  await setBlob(app, u1.token.token, workspace.id, buffer2);

  const size = await getWorkspaceBlobsSize(app, u1.token.token, workspace.id);
  t.is(size, 4, 'failed to collect blob sizes');
});

test('should calc all blobs size', async t => {
  const u1 = await signUp(app, 'u1', 'u1@affine.pro', '1');

  const workspace1 = await createWorkspace(app, u1.token.token);

  const buffer1 = Buffer.from([0, 0]);
  await setBlob(app, u1.token.token, workspace1.id, buffer1);
  const buffer2 = Buffer.from([0, 1]);
  await setBlob(app, u1.token.token, workspace1.id, buffer2);

  const workspace2 = await createWorkspace(app, u1.token.token);

  const buffer3 = Buffer.from([0, 0]);
  await setBlob(app, u1.token.token, workspace2.id, buffer3);
  const buffer4 = Buffer.from([0, 1]);
  await setBlob(app, u1.token.token, workspace2.id, buffer4);

  const size = await collectAllBlobSizes(app, u1.token.token);
  t.is(size, 8, 'failed to collect all blob sizes');

  const size1 = await checkBlobSize(
    app,
    u1.token.token,
    workspace1.id,
    10 * 1024 * 1024 * 1024 - 8
  );
  t.is(size1, 0, 'failed to check blob size');

  const size2 = await checkBlobSize(
    app,
    u1.token.token,
    workspace1.id,
    10 * 1024 * 1024 * 1024 - 7
  );
  t.is(size2, -1, 'failed to check blob size');
});

test('should be able calc quota after switch plan', async t => {
  const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1');

  const workspace1 = await createWorkspace(app, u1.token.token);

  const buffer1 = Buffer.from([0, 0]);
  await setBlob(app, u1.token.token, workspace1.id, buffer1);
  const buffer2 = Buffer.from([0, 1]);
  await setBlob(app, u1.token.token, workspace1.id, buffer2);

  const workspace2 = await createWorkspace(app, u1.token.token);

  const buffer3 = Buffer.from([0, 0]);
  await setBlob(app, u1.token.token, workspace2.id, buffer3);
  const buffer4 = Buffer.from([0, 1]);
  await setBlob(app, u1.token.token, workspace2.id, buffer4);

  const size1 = await checkBlobSize(
    app,
    u1.token.token,
    workspace1.id,
    10 * 1024 * 1024 * 1024 - 8
  );
  t.is(size1, 0, 'failed to check free plan blob size');

  quota.switchUserQuota(u1.id, QuotaType.Quota_ProPlanV1);

  const size2 = await checkBlobSize(
    app,
    u1.token.token,
    workspace1.id,
    100 * 1024 * 1024 * 1024 - 8
  );
  t.is(size2, 0, 'failed to check pro plan blob size');
});