AFFiNE-Mirror/packages/frontend/apps/electron/test/main/byok-storage.spec.ts

import path from 'node:path';

import fs from 'fs-extra';
import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest';

const tmpDir = path.join(__dirname, 'tmp-byok-storage');
let disposeWorkspaceByokStorage: (() => void) | undefined;

vi.mock('electron', () => ({
  app: {
    getPath: () => tmpDir,
    on: vi.fn(),
  },
  safeStorage: {
    isEncryptionAvailable: () => true,
    encryptString: (value: string) => Buffer.from(value, 'utf-8'),
    decryptString: (value: Buffer) => value.toString('utf-8'),
  },
}));

beforeEach(async () => {
  vi.resetModules();
  disposeWorkspaceByokStorage = undefined;
  await fs.remove(tmpDir);
});

afterEach(async () => {
  disposeWorkspaceByokStorage?.();
  vi.resetModules();
  await fs.remove(tmpDir);
});

describe('byok storage handlers', () => {
  test('stores encrypted local keys and keeps lease providers sorted', async () => {
    const { byokStorageHandlers, disposeWorkspaceByokStorage: dispose } =
      await import('@affine/electron/main/byok-storage/handlers');
    disposeWorkspaceByokStorage = dispose;
    const ipcEvent = undefined;

    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
      id: 'local-openai',
      provider: 'openai',
      name: 'OpenAI',
      apiKey: 'sk-openai',
      sortOrder: 1,
    });
    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
      id: 'local-gemini',
      provider: 'gemini',
      name: 'Gemini',
      apiKey: 'sk-gemini',
      sortOrder: 0,
    });

    const list = await byokStorageHandlers.listWorkspaceKeys(
      ipcEvent,
      'workspace-1'
    );
    expect(list.map(key => key.id)).toEqual(['local-gemini', 'local-openai']);
    expect(JSON.stringify(list)).not.toContain('sk-openai');

    const reordered = await byokStorageHandlers.reorderWorkspaceKeys(
      ipcEvent,
      'workspace-1',
      ['local-openai', 'local-gemini']
    );
    expect(reordered.map(key => key.id)).toEqual([
      'local-openai',
      'local-gemini',
    ]);

    const leaseProviders = await byokStorageHandlers.getWorkspaceLeaseProviders(
      ipcEvent,
      'workspace-1'
    );
    expect(leaseProviders.map(key => key.apiKey)).toEqual([
      'sk-openai',
      'sk-gemini',
    ]);

    await byokStorageHandlers.clearWorkspaceKeys(ipcEvent, 'workspace-1');
    await expect(
      byokStorageHandlers.listWorkspaceKeys(ipcEvent, 'workspace-1')
    ).resolves.toEqual([]);
  });

  test('preserves existing local key fields during partial updates', async () => {
    const { byokStorageHandlers, disposeWorkspaceByokStorage: dispose } =
      await import('@affine/electron/main/byok-storage/handlers');
    disposeWorkspaceByokStorage = dispose;
    const ipcEvent = undefined;

    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
      id: 'local-openai',
      provider: 'openai',
      name: 'OpenAI',
      description: 'Primary key',
      apiKey: 'sk-openai',
      endpoint: 'https://api.openai.example/v1',
      sortOrder: 4,
      enabled: false,
    });

    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
      id: 'local-openai',
      provider: 'openai',
      name: 'OpenAI renamed',
      apiKey: 'sk-openai-next',
    });

    const [publicKey] = await byokStorageHandlers.listWorkspaceKeys(
      ipcEvent,
      'workspace-1'
    );
    expect(publicKey).toMatchObject({
      id: 'local-openai',
      name: 'OpenAI renamed',
      description: 'Primary key',
      endpoint: 'https://api.openai.example/v1',
      sortOrder: 4,
      enabled: false,
    });

    const [leaseProvider] =
      await byokStorageHandlers.getWorkspaceLeaseProviders(
        ipcEvent,
        'workspace-1'
      );
    expect(leaseProvider).toBeUndefined();

    await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
      id: 'local-openai',
      provider: 'openai',
      name: 'OpenAI renamed again',
      enabled: true,
    });

    const [enabledLeaseProvider] =
      await byokStorageHandlers.getWorkspaceLeaseProviders(
        ipcEvent,
        'workspace-1'
      );
    expect(enabledLeaseProvider).toMatchObject({
      name: 'OpenAI renamed again',
      apiKey: 'sk-openai-next',
      endpoint: 'https://api.openai.example/v1',
      sortOrder: 4,
      enabled: true,
    });
  });
});