zn-ai/tests/uv-setup.test.ts

// @vitest-environment node

import { beforeEach, describe, expect, it, vi } from 'vitest';

const childProcessMocks = vi.hoisted(() => ({
  execSync: vi.fn(),
  spawn: vi.fn(),
}));

const fsMocks = vi.hoisted(() => ({
  existsSync: vi.fn(),
}));

const electronMocks = vi.hoisted(() => ({
  app: {
    isPackaged: false,
  },
}));

const uvEnvMocks = vi.hoisted(() => ({
  getUvMirrorEnv: vi.fn(async () => ({})),
}));

const loggerMocks = vi.hoisted(() => ({
  info: vi.fn(),
  warn: vi.fn(),
  debug: vi.fn(),
  error: vi.fn(),
}));

vi.mock('node:child_process', () => ({
  execSync: childProcessMocks.execSync,
  spawn: childProcessMocks.spawn,
}));

vi.mock('node:fs', () => ({
  existsSync: fsMocks.existsSync,
}));

vi.mock('node:path', () => ({
  join: (...parts: string[]) => parts.join('/'),
}));

vi.mock('electron', () => ({
  app: electronMocks.app,
}));

vi.mock('../electron/utils/uv-env', () => ({
  getUvMirrorEnv: uvEnvMocks.getUvMirrorEnv,
}));

vi.mock('@electron/service/logger', () => ({
  default: loggerMocks,
}));

import {
  isPythonReady,
  setupManagedPython,
} from '../electron/utils/uv-setup';

function createSpawnChild(options: {
  closeCode?: number;
  error?: Error;
  stdout?: string[];
  stderr?: string[];
}) {
  const childListeners = new Map<string, Array<(value?: unknown) => void>>();
  const stdoutListeners = new Map<string, Array<(value: unknown) => void>>();
  const stderrListeners = new Map<string, Array<(value: unknown) => void>>();

  const emitAll = (listeners: Map<string, Array<(value?: unknown) => void>>, event: string, value?: unknown) => {
    for (const listener of listeners.get(event) || []) {
      listener(value);
    }
  };

  const child = {
    stdout: {
      on(event: string, listener: (value: unknown) => void) {
        const next = stdoutListeners.get(event) || [];
        next.push(listener);
        stdoutListeners.set(event, next);
      },
    },
    stderr: {
      on(event: string, listener: (value: unknown) => void) {
        const next = stderrListeners.get(event) || [];
        next.push(listener);
        stderrListeners.set(event, next);
      },
    },
    on(event: string, listener: (value?: unknown) => void) {
      const next = childListeners.get(event) || [];
      next.push(listener);
      childListeners.set(event, next);
      return child;
    },
  };

  queueMicrotask(() => {
    for (const chunk of options.stdout || []) {
      emitAll(stdoutListeners as Map<string, Array<(value?: unknown) => void>>, 'data', Buffer.from(chunk));
    }
    for (const chunk of options.stderr || []) {
      emitAll(stderrListeners as Map<string, Array<(value?: unknown) => void>>, 'data', Buffer.from(chunk));
    }

    if (options.error) {
      emitAll(childListeners, 'error', options.error);
      return;
    }

    emitAll(childListeners, 'close', options.closeCode ?? 0);
  });

  return child;
}

describe('uv setup', () => {
  beforeEach(() => {
    childProcessMocks.execSync.mockReset();
    childProcessMocks.spawn.mockReset();
    fsMocks.existsSync.mockReset();
    uvEnvMocks.getUvMirrorEnv.mockReset();
    uvEnvMocks.getUvMirrorEnv.mockResolvedValue({});
    loggerMocks.info.mockReset();
    loggerMocks.warn.mockReset();
    loggerMocks.debug.mockReset();
    loggerMocks.error.mockReset();
    electronMocks.app.isPackaged = false;
  });

  it('does not spawn uv when bundled and PATH uv are both missing', async () => {
    fsMocks.existsSync.mockReturnValue(false);
    childProcessMocks.execSync.mockImplementation(() => {
      throw new Error('uv not found');
    });

    await expect(isPythonReady()).resolves.toBe(false);
    await expect(setupManagedPython()).rejects.toThrow('uv is required for managed Python setup but is unavailable');

    expect(childProcessMocks.spawn).not.toHaveBeenCalled();
  });

  it('does not fall back to literal uv when PATH lookup returns empty output', async () => {
    fsMocks.existsSync.mockReturnValue(false);
    childProcessMocks.execSync.mockReturnValue('\r\n');

    await expect(isPythonReady()).resolves.toBe(false);
    await expect(setupManagedPython()).rejects.toThrow('uv is required for managed Python setup but is unavailable');

    expect(childProcessMocks.spawn).not.toHaveBeenCalled();
  });

  it('retries Python install without mirror when bundled uv exists', async () => {
    fsMocks.existsSync.mockImplementation((value: unknown) => String(value).endsWith('uv.exe'));
    childProcessMocks.execSync.mockImplementation(() => {
      throw new Error('uv not on PATH');
    });
    uvEnvMocks.getUvMirrorEnv.mockResolvedValue({
      UV_INDEX_URL: 'https://mirror.example/simple',
    });

    childProcessMocks.spawn
      .mockImplementationOnce(() => createSpawnChild({
        closeCode: 1,
        stderr: ['mirror failed'],
      }))
      .mockImplementationOnce(() => createSpawnChild({
        closeCode: 0,
        stdout: ['installed'],
      }))
      .mockImplementationOnce(() => createSpawnChild({
        closeCode: 0,
        stdout: ['C:\\Python312\\python.exe'],
      }));

    await expect(setupManagedPython()).resolves.toBeUndefined();

    expect(childProcessMocks.spawn).toHaveBeenCalledTimes(3);
    expect(String(childProcessMocks.spawn.mock.calls[0]?.[0] || '')).toContain('uv.exe');
    expect(childProcessMocks.spawn.mock.calls[0]?.[2]?.env?.UV_INDEX_URL).toBe('https://mirror.example/simple');
    expect(childProcessMocks.spawn.mock.calls[1]?.[2]?.env?.UV_INDEX_URL).toBeUndefined();
    expect(loggerMocks.warn).toHaveBeenCalledWith('Python install attempt 1 failed:', expect.any(Error));
    expect(loggerMocks.info).toHaveBeenCalledWith('Retrying Python install without mirror...');
  });
});