inkling/tests/unit/AiWorker.test.ts

import { describe, it, expect, beforeEach, vi } from 'vitest';
import Database from 'better-sqlite3';
import { runMigrations } from '@main/db/migrations/index.js';
import { NoteRepository } from '@main/repository/NoteRepository.js';
import { AiWorker } from '@main/ai/AiWorker.js';
import type { InferenceProvider } from '@main/ai/InferenceProvider.js';
import type { AiResponse } from '@main/ai/schema.js';

function makeProvider(overrides: Partial<InferenceProvider> = {}): InferenceProvider {
  return {
    name: 'mock',
    generate: vi.fn(async (): Promise<AiResponse> => ({
      title: '제목', summary: 'a\nb\nc', tags: ['tag'], dueDate: null
    })),
    healthCheck: vi.fn(async () => ({ ok: true })),
    ...overrides
  } as InferenceProvider;
}

describe('AiWorker', () => {
  let db: Database.Database;
  let repo: NoteRepository;

  beforeEach(() => {
    db = new Database(':memory:');
    runMigrations(db);
    repo = new NoteRepository(db);
  });

  it('processes a pending job and marks done', async () => {
    const { id } = repo.create({ rawText: 'x' });
    const updates: string[] = [];
    const w = new AiWorker(repo, makeProvider(), {
      backoffsMs: [0, 0, 0],
      onUpdate: (note) => updates.push(note.aiStatus)
    });
    await w.enqueue(id);
    await w.drain();
    expect(repo.findById(id)?.aiStatus).toBe('done');
    expect(updates).toContain('done');
  });

  it('retries 3 times then marks failed', async () => {
    const { id } = repo.create({ rawText: 'x' });
    const provider = makeProvider({
      generate: vi.fn(async () => { throw new Error('boom'); })
    });
    const w = new AiWorker(repo, provider, { backoffsMs: [0, 0, 0] });
    await w.enqueue(id);
    await w.drain();
    const note = repo.findById(id)!;
    expect(note.aiStatus).toBe('failed');
    expect(note.aiError).toContain('boom');
    expect(provider.generate).toHaveBeenCalledTimes(3);
  });

  it('loadFromDb re-queues all pending', async () => {
    const a = repo.create({ rawText: 'a' }).id;
    const b = repo.create({ rawText: 'b' }).id;
    const w = new AiWorker(repo, makeProvider(), { backoffsMs: [0, 0, 0] });
    await w.loadFromDb();
    await w.drain();
    expect(repo.findById(a)?.aiStatus).toBe('done');
    expect(repo.findById(b)?.aiStatus).toBe('done');
  });

  it('processes sequentially (concurrency 1)', async () => {
    const ids = [repo.create({ rawText: 'a' }).id, repo.create({ rawText: 'b' }).id];
    let running = 0;
    let max = 0;
    const provider = makeProvider({
      generate: vi.fn(async () => {
        running++; max = Math.max(max, running);
        await new Promise((r) => setTimeout(r, 10));
        running--;
        return { title: '제목', summary: 'a\nb\nc', tags: [], dueDate: null };
      })
    });
    const w = new AiWorker(repo, provider, { backoffsMs: [0, 0, 0] });
    for (const id of ids) await w.enqueue(id);
    await w.drain();
    expect(max).toBe(1);
  });

  it('rule parser match takes priority over AI dueDate', async () => {
    const provider = {
      name: 'mock',
      generate: async (_input: any) => ({
        title: '내일',
        summary: 'a\nb\nc',
        tags: [],
        dueDate: '2026-12-31' // AI returns far-future
      }),
      healthCheck: async () => ({ ok: true })
    } as any;
    const w = new AiWorker(repo, provider, {
      backoffsMs: [0],
      now: () => new Date('2026-04-26T00:00:00.000Z')
    });
    const { id } = repo.create({ rawText: '내일 회의' });
    await w.enqueue(id);
    await w.drain();
    const note = repo.findById(id)!;
    expect(note.dueDate).toBe('2026-04-27'); // 내일 from rule, not AI's 12-31
  });

  it('rule null + AI value → AI used', async () => {
    const provider = {
      name: 'mock',
      generate: async () => ({
        title: '월말 마감',
        summary: 'a\nb\nc',
        tags: [],
        dueDate: '2026-04-30' // AI resolves "월말"
      }),
      healthCheck: async () => ({ ok: true })
    } as any;
    const w = new AiWorker(repo, provider, {
      backoffsMs: [0],
      now: () => new Date('2026-04-26T00:00:00.000Z')
    });
    const { id } = repo.create({ rawText: '월말 마감' });
    await w.enqueue(id);
    await w.drain();
    const note = repo.findById(id)!;
    expect(note.dueDate).toBe('2026-04-30');
  });

  it('rule null + AI null → null', async () => {
    const provider = {
      name: 'mock',
      generate: async () => ({
        title: '아무 메모',
        summary: 'a\nb\nc',
        tags: [],
        dueDate: null
      }),
      healthCheck: async () => ({ ok: true })
    } as any;
    const w = new AiWorker(repo, provider, {
      backoffsMs: [0],
      now: () => new Date('2026-04-26T00:00:00.000Z')
    });
    const { id } = repo.create({ rawText: '아무 메모' });
    await w.enqueue(id);
    await w.drain();
    const note = repo.findById(id)!;
    expect(note.dueDate).toBeNull();
  });

  it('passes todayKst to provider.generate', async () => {
    const seen: any = {};
    const provider = {
      name: 'mock',
      generate: async (input: any) => {
        seen.todayKst = input.todayKst;
        return { title: '메모', summary: 'a\nb\nc', tags: [], dueDate: null };
      },
      healthCheck: async () => ({ ok: true })
    } as any;
    const w = new AiWorker(repo, provider, {
      backoffsMs: [0],
      now: () => new Date('2026-04-26T15:00:00.000Z') // 04-27 00:00 KST
    });
    const { id } = repo.create({ rawText: 'x' });
    await w.enqueue(id);
    await w.drain();
    expect(seen.todayKst).toBe('2026-04-27');
  });
});