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

import os from 'node:os';
import path from 'node:path';

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

const electronMock = vi.hoisted(() => ({
  tmpDir: '',
  appOn: vi.fn(),
  isEncryptionAvailable: vi.fn(() => true),
  encryptString: vi.fn((value: string) => Buffer.from(value, 'utf-8')),
  decryptString: vi.fn((value: Buffer) => value.toString('utf-8')),
}));

let disposeWorkspaceByokStorage: (() => void) | undefined;

vi.mock('electron', () => ({
  app: {
    getPath: () => electronMock.tmpDir,
    on: electronMock.appOn,
  },
  safeStorage: {
    isEncryptionAvailable: electronMock.isEncryptionAvailable,
    encryptString: electronMock.encryptString,
    decryptString: electronMock.decryptString,
  },
}));

vi.mock('../../src/main/logger', () => ({
  logger: {
    error: vi.fn(),
  },
}));

beforeEach(async () => {
  vi.useRealTimers();
  vi.resetModules();
  electronMock.appOn.mockReset();
  electronMock.isEncryptionAvailable.mockReset().mockReturnValue(true);
  electronMock.encryptString
    .mockReset()
    .mockImplementation((value: string) => Buffer.from(value, 'utf-8'));
  electronMock.decryptString
    .mockReset()
    .mockImplementation((value: Buffer) => value.toString('utf-8'));
  disposeWorkspaceByokStorage = undefined;
  electronMock.tmpDir = await fs.mkdtemp(
    path.join(os.tmpdir(), 'affine-byok-storage-')
  );
});

afterEach(async () => {
  disposeWorkspaceByokStorage?.();
  disposeWorkspaceByokStorage = undefined;
  vi.resetModules();
  if (electronMock.tmpDir) {
    await fs.remove(electronMock.tmpDir);
  }
  electronMock.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('does not write local keys when secure storage is unavailable', async () => {
    electronMock.isEncryptionAvailable.mockReturnValue(false);

    const { byokStorageHandlers, disposeWorkspaceByokStorage: dispose } =
      await import('@affine/electron/main/byok-storage/handlers');
    disposeWorkspaceByokStorage = dispose;
    const ipcEvent = undefined;

    await expect(byokStorageHandlers.isSupported()).resolves.toBe(false);
    await expect(
      byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', {
        id: 'local-openai',
        provider: 'openai',
        name: 'OpenAI',
        apiKey: 'sk-openai',
      })
    ).rejects.toThrow('Secure BYOK key storage is not available.');
    expect(electronMock.encryptString).not.toHaveBeenCalled();
  });

  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,
    });
  });
});