diff --git a/packages/workspace/src/affine/__tests__/api.spec.ts b/packages/workspace/src/affine/__tests__/api.spec.ts index 297a37ec2f..bd69a42ea7 100644 --- a/packages/workspace/src/affine/__tests__/api.spec.ts +++ b/packages/workspace/src/affine/__tests__/api.spec.ts @@ -15,9 +15,11 @@ import { faker } from '@faker-js/faker'; import { beforeEach, describe, expect, test, vi } from 'vitest'; import { + createUserApis, createWorkspaceApis, createWorkspaceResponseSchema, RequestError, + usageResponseSchema, } from '../api'; import { createAffineAuth, @@ -27,6 +29,7 @@ import { } from '../login'; let workspaceApis: ReturnType; +let userApis: ReturnType; let affineAuth: ReturnType; let statusApis: ReturnType; @@ -45,6 +48,7 @@ beforeEach(() => { beforeEach(() => { affineAuth = createAffineAuth('http://localhost:3000/'); + userApis = createUserApis('http://localhost:3000/'); workspaceApis = createWorkspaceApis('http://localhost:3000/'); statusApis = createStatusApis('http://localhost:3000/'); }); @@ -203,13 +207,14 @@ describe('api', () => { const binary = Workspace.Y.encodeStateAsUpdate(workspace.doc); const data = await workspaceApis.createWorkspace(new Blob([binary])); createWorkspaceResponseSchema.parse(data); + const workspaceId = data.id; const blobId = await workspaceApis.uploadBlob( - data.id, + workspaceId, imageBuffer, 'image/png' ); expect(blobId).toBeTypeOf('string'); - const arrayBuffer = await workspaceApis.getBlob(blobId); + const arrayBuffer = await workspaceApis.getBlob(workspaceId, blobId); expect(arrayBuffer).toBeInstanceOf(ArrayBuffer); expect(arrayBuffer.byteLength).toEqual(imageBuffer.byteLength); expect(Buffer.from(arrayBuffer)).toEqual(imageBuffer); @@ -237,4 +242,31 @@ describe('api', () => { timeout: 30000, } ); + + test( + 'usage', + async () => { + const usageResponse = await userApis.getUsage(); + usageResponseSchema.parse(usageResponse); + const id = await createWorkspace(workspaceApis); + const path = require.resolve('@affine-test/fixtures/smile.png'); + const imageBuffer = await readFile(path); + const blobId = await workspaceApis.uploadBlob( + id, + imageBuffer, + 'image/png' + ); + const buffer = await workspaceApis.getBlob(id, blobId); + expect(buffer.byteLength).toEqual(imageBuffer.byteLength); + const newUsageResponse = await userApis.getUsage(); + usageResponseSchema.parse(newUsageResponse); + expect(usageResponse.blob_usage.usage).not.equals( + newUsageResponse.blob_usage.usage + ); + expect(newUsageResponse.blob_usage.usage).equals(96); + }, + { + timeout: 30000, + } + ); }); diff --git a/packages/workspace/src/affine/api/index.ts b/packages/workspace/src/affine/api/index.ts index 1c8437e2ff..f4a5a4a8c7 100644 --- a/packages/workspace/src/affine/api/index.ts +++ b/packages/workspace/src/affine/api/index.ts @@ -39,8 +39,27 @@ export interface GetUserByEmailParams { workspace_id: string; } +export const usageResponseSchema = z.object({ + blob_usage: z.object({ + usage: z.number(), + max_usage: z.number(), + }), +}); + +export type UsageResponse = z.infer; + export function createUserApis(prefixUrl = '/') { return { + getUsage: async (): Promise => { + const auth = getLoginStorage(); + assertExists(auth); + return fetch(prefixUrl + 'api/resource/usage', { + method: 'GET', + headers: { + Authorization: auth.token, + }, + }).then(r => r.json()); + }, getUserByEmail: async ( params: GetUserByEmailParams ): Promise => { @@ -309,7 +328,7 @@ export function createWorkspaceApis(prefixUrl = '/') { if (mb > 10) { throw new RequestError(MessageCode.blobTooLarge); } - return fetch(prefixUrl + 'api/blob', { + return fetch(prefixUrl + `api/workspace/${workspaceId}/blob`, { method: 'PUT', body: arrayBuffer, headers: { @@ -318,10 +337,13 @@ export function createWorkspaceApis(prefixUrl = '/') { }, }).then(r => r.text()); }, - getBlob: async (blobId: string): Promise => { + getBlob: async ( + workspaceId: string, + blobId: string + ): Promise => { const auth = getLoginStorage(); assertExists(auth); - return fetch(prefixUrl + `api/blob/${blobId}`, { + return fetch(prefixUrl + `api/workspace/${workspaceId}/blob/${blobId}`, { method: 'GET', headers: { Authorization: auth.token,