badgr-cli 1.1.0 → 1.1.2

package/tests/template.test.js

@@ -1,556 +0,0 @@
-/**
- * Template catalog + dispatch tests
- *
- * Covers:
- *   - Catalog integrity: all 10 templates, required fields, uniqueness
- *   - buildTemplateFlags: correct flags for each template type
- *   - parseTemplateOverrides: flag parsing, user overrides win
- *   - badgr serve template <name>: dispatches via serveCommand (mocked API)
- *   - badgr run template <name>: dispatches via runCommand (mocked API)
- *   - Type-mismatch routing guard (job via serve, endpoint via run)
- *   - Unknown template guard
- */
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import {
-  TEMPLATES,
-  TEMPLATE_MAP,
-  buildTemplateFlags,
-  parseTemplateOverrides,
-} from '../src/catalog.js';
-import { serveCommand } from '../src/commands/serve.js';
-import { runCommand }   from '../src/commands/run.js';
-
-// ── Module mocks ──────────────────────────────────────────────────────────────
-
-vi.mock('../src/api.js', () => ({
-  callApi:             vi.fn(),
-  terminateDeployment: vi.fn().mockResolvedValue({}),
-  listDeployments:     vi.fn().mockResolvedValue({ deployments: [], count: 0 }),
-}));
-
-vi.mock('../src/store.js', () => ({
-  selectedComputeFromDeployment: (dep) => ({
-    gpu: dep.gpu_type ?? null,
-    gpuCount: dep.gpu_count ?? null,
-    vcpus: dep.selected_vcpus ?? null,
-    ramGb: dep.selected_ram_gb ?? null,
-    vramGb: dep.selected_vram_gb ?? null,
-  }),
-  addDeployment:        vi.fn(),
-  addReceipt:           vi.fn(),
-  updateReceipt:        vi.fn(),
-  generateReceiptId:    vi.fn(() => 'rcpt-tmpl-001'),
-  generateDeploymentId: vi.fn(() => 'dep-tmpl-001'),
-  listDeployments:      vi.fn(() => []),
-  listReceipts:         vi.fn(() => []),
-  findDeployment:       vi.fn(() => null),
-  updateDeployment:     vi.fn(),
-  removeDeployment:     vi.fn(),
-}));
-
-import * as api   from '../src/api.js';
-import * as store from '../src/store.js';
-
-// ── Shared helpers ────────────────────────────────────────────────────────────
-
-const config = { apiKey: 'sk-test', baseUrl: 'https://api.test/v1' };
-
-const chalk = {
-  bold: s => s, dim: s => s, red: s => s,
-  yellow: s => s, green: s => s, cyan: s => s,
-};
-
-function makeServeDep(overrides = {}) {
-  return {
-    deployment_id: 'dep-tmpl-001',
-    status:        'running',
-    gpu_type:      'RTX_4090',
-    gpu_count:     1,
-    cost_per_hour: 0.50,
-    provider:      'runpod',
-    receipt_id:    'rcpt-tmpl-001',
-    tier:          '1',
-    endpoint_url:  'https://dep-tmpl-001.aibadgr.com',
-    ...overrides,
-  };
-}
-
-function makeRunDep(overrides = {}) {
-  return {
-    deployment_id: 'dep-tmpl-001',
-    status:        'running',
-    gpu_type:      'A100',
-    gpu_count:     1,
-    cost_per_hour: 2.50,
-    provider:      'runpod',
-    receipt_id:    'rcpt-tmpl-001',
-    tier:          '1',
-    ...overrides,
-  };
-}
-
-beforeEach(() => {
-  vi.useFakeTimers();
-  process.exitCode = undefined;
-  vi.spyOn(console, 'log').mockImplementation(() => {});
-  vi.spyOn(console, 'error').mockImplementation(() => {});
-  vi.spyOn(process.stdout, 'write').mockImplementation(() => true);
-  vi.clearAllMocks();
-  store.generateReceiptId.mockReturnValue('rcpt-tmpl-001');
-  api.terminateDeployment.mockResolvedValue({});
-});
-
-afterEach(() => {
-  vi.useRealTimers();
-  vi.restoreAllMocks();
-  process.exitCode = undefined;
-});
-
-// ─────────────────────────────────────────────────────────────────────────────
-// 1. Catalog integrity
-// ─────────────────────────────────────────────────────────────────────────────
-
-describe('TEMPLATES catalog', () => {
-  const EXPECTED_NAMES = [
-    'comfyui', 'axolotl', 'batch-inference', 'unsloth', 'vllm', 'openwebui', 'llama-cpp',
-    'invokeai', 'kohya-ss', 'text-gen-webui', 'sglang', 'tgi',
-    'auto1111', 'forge', 'nerfstudio', 'openfold', 'blender-render',
-    'openmm', 'gromacs', 'lammps', 'diffusers', 'torchtune',
-  ];
-
-  it('contains exactly 22 templates', () => {
-    expect(TEMPLATES).toHaveLength(22);
-  });
-
-  it('contains all expected template names', () => {
-    const names = TEMPLATES.map(t => t.name);
-    for (const name of EXPECTED_NAMES) {
-      expect(names, `missing: ${name}`).toContain(name);
-    }
-  });
-
-  it('template names are unique', () => {
-    const names = TEMPLATES.map(t => t.name);
-    expect(new Set(names).size).toBe(names.length);
-  });
-
-  it('every template has required fields', () => {
-    for (const t of TEMPLATES) {
-      expect(t.name,        `${t.name}: name`).toBeTruthy();
-      expect(t.title,       `${t.name}: title`).toBeTruthy();
-      expect(t.description, `${t.name}: description`).toBeTruthy();
-      expect(['endpoint', 'job'], `${t.name}: type`).toContain(t.type);
-      expect(t.image,       `${t.name}: image`).toBeTruthy();
-      expect(t.gpu,         `${t.name}: gpu`).toBeTruthy();
-      expect(t.gpu_count,   `${t.name}: gpu_count`).toBeGreaterThan(0);
-      expect(t.min_vram_gb, `${t.name}: min_vram_gb`).toBeGreaterThan(0);
-    }
-  });
-
-  it('endpoint templates have port and health_path', () => {
-    for (const t of TEMPLATES.filter(t => t.type === 'endpoint')) {
-      expect(t.port,        `${t.name}: port`).toBeGreaterThan(0);
-      expect(t.health_path, `${t.name}: health_path`).toBeTruthy();
-    }
-  });
-
-  it('job templates do not have a port', () => {
-    for (const t of TEMPLATES.filter(t => t.type === 'job')) {
-      expect(t.port, `${t.name}: port should be absent`).toBeFalsy();
-    }
-  });
-
-  it('no env value is undefined', () => {
-    for (const t of TEMPLATES) {
-      for (const [k, v] of Object.entries(t.env ?? {})) {
-        expect(v, `${t.name}.env.${k}`).not.toBeUndefined();
-      }
-    }
-  });
-
-  it('vLLM requires ≥ 24 GB VRAM', () => {
-    expect(TEMPLATE_MAP['vllm'].min_vram_gb).toBeGreaterThanOrEqual(24);
-  });
-
-  it('SGLang requires ≥ 40 GB VRAM', () => {
-    expect(TEMPLATE_MAP['sglang'].min_vram_gb).toBeGreaterThanOrEqual(40);
-  });
-
-  it('job templates use high-VRAM GPUs', () => {
-    const highVram = ['A100', 'H100', 'L40S', 'A6000', 'RTX_4090'];
-    for (const t of TEMPLATES.filter(t => t.type === 'job')) {
-      expect(highVram, `${t.name}: gpu`).toContain(t.gpu);
-    }
-  });
-
-  it('TEMPLATE_MAP keys match template names', () => {
-    for (const [k, t] of Object.entries(TEMPLATE_MAP)) {
-      expect(k).toBe(t.name);
-    }
-    expect(Object.keys(TEMPLATE_MAP)).toHaveLength(22);
-  });
-});
-
-// ─────────────────────────────────────────────────────────────────────────────
-// 2. buildTemplateFlags
-// ─────────────────────────────────────────────────────────────────────────────
-
-describe('buildTemplateFlags', () => {
-  it('produces --image and --gpu for every template', () => {
-    for (const t of TEMPLATES) {
-      const flags = buildTemplateFlags(t, {});
-      expect(flags).toContain('--image');
-      expect(flags).toContain(t.image);
-      expect(flags).toContain('--gpu');
-      expect(flags).toContain(t.gpu);
-    }
-  });
-
-  it('endpoint templates include --health-path', () => {
-    for (const t of TEMPLATES.filter(t => t.type === 'endpoint')) {
-      const flags = buildTemplateFlags(t, {});
-      expect(flags, `${t.name}: --health-path`).toContain('--health-path');
-      expect(flags, `${t.name}: health_path value`).toContain(t.health_path);
-    }
-  });
-
-  it('job templates do not include --health-path', () => {
-    for (const t of TEMPLATES.filter(t => t.type === 'job')) {
-      const flags = buildTemplateFlags(t, {});
-      expect(flags, `${t.name}: no --health-path`).not.toContain('--health-path');
-    }
-  });
-
-  it('strips placeholder env values (those starting with <)', () => {
-    for (const t of TEMPLATES) {
-      const flags = buildTemplateFlags(t, {});
-      const envValues = flags
-        .filter((_, i) => i > 0 && flags[i - 1] === '--env')
-        .map(kv => kv.split('=').slice(1).join('='));
-      for (const v of envValues) {
-        expect(v, `${t.name}: placeholder passed through`).not.toMatch(/^</);
-      }
-    }
-  });
-
-  it('vllm includes MODEL, MAX_MODEL_LEN, TENSOR_PARALLEL_SIZE env vars', () => {
-    const flags = buildTemplateFlags(TEMPLATE_MAP['vllm'], {});
-    expect(flags).toContain('MODEL=meta-llama/Llama-3.1-8B-Instruct');
-    expect(flags.join(' ')).toContain('MAX_MODEL_LEN=8192');
-    expect(flags.join(' ')).toContain('TENSOR_PARALLEL_SIZE=1');
-  });
-
-  it('llama-cpp includes LLAMA_ARG_HF_REPO and LLAMA_ARG_HF_FILE', () => {
-    const flags = buildTemplateFlags(TEMPLATE_MAP['llama-cpp'], {});
-    expect(flags.join(' ')).toContain('LLAMA_ARG_HF_REPO=');
-    expect(flags.join(' ')).toContain('LLAMA_ARG_HF_FILE=');
-  });
-
-  it('user --gpu override replaces template default', () => {
-    const flags = buildTemplateFlags(TEMPLATE_MAP['vllm'], { gpu: 'A100' });
-    const gpuIdx = flags.indexOf('--gpu');
-    expect(flags[gpuIdx + 1]).toBe('A100');
-  });
-
-  it('user --env override replaces template env default', () => {
-    const flags = buildTemplateFlags(TEMPLATE_MAP['vllm'], {
-      env: { MODEL: 'mistralai/Mistral-7B-v0.1' },
-    });
-    expect(flags).toContain('MODEL=mistralai/Mistral-7B-v0.1');
-    expect(flags).not.toContain('MODEL=meta-llama/Llama-3.1-8B-Instruct');
-  });
-
-  it('--max-cost is included when supplied in overrides', () => {
-    const flags = buildTemplateFlags(TEMPLATE_MAP['comfyui'], { maxCost: 5 });
-    expect(flags).toContain('--max-cost');
-    expect(flags).toContain('5');
-  });
-
-  it('--count is included only when > 1', () => {
-    const single = buildTemplateFlags(TEMPLATE_MAP['vllm'], {});
-    expect(single).not.toContain('--count');
-
-    const multi = buildTemplateFlags(TEMPLATE_MAP['vllm'], { count: 2 });
-    expect(multi).toContain('--count');
-    expect(multi).toContain('2');
-  });
-
-  it('--persistent flag is forwarded', () => {
-    const flags = buildTemplateFlags(TEMPLATE_MAP['comfyui'], { persistent: true });
-    expect(flags).toContain('--persistent');
-  });
-});
-
-// ─────────────────────────────────────────────────────────────────────────────
-// 3. parseTemplateOverrides
-// ─────────────────────────────────────────────────────────────────────────────
-
-describe('parseTemplateOverrides', () => {
-  it('parses --max-cost', () => {
-    const o = parseTemplateOverrides(['--max-cost', '10']);
-    expect(o.maxCost).toBe(10);
-  });
-
-  it('parses --gpu', () => {
-    const o = parseTemplateOverrides(['--gpu', 'A100']);
-    expect(o.gpu).toBe('A100');
-  });
-
-  it('parses --max-runtime', () => {
-    const o = parseTemplateOverrides(['--max-runtime', '120']);
-    expect(o.maxRuntime).toBe(120);
-  });
-
-  it('parses multiple --env flags', () => {
-    const o = parseTemplateOverrides(['--env', 'MODEL=mistral', '--env', 'HF_TOKEN=abc123']);
-    expect(o.env.MODEL).toBe('mistral');
-    expect(o.env.HF_TOKEN).toBe('abc123');
-  });
-
-  it('parses --persistent', () => {
-    const o = parseTemplateOverrides(['--persistent']);
-    expect(o.persistent).toBe(true);
-  });
-
-  it('parses --no-wait', () => {
-    const o = parseTemplateOverrides(['--no-wait']);
-    expect(o.noWait).toBe(true);
-  });
-
-  it('ignores unknown flags without throwing', () => {
-    expect(() => parseTemplateOverrides(['--unknown-flag', 'value'])).not.toThrow();
-  });
-
-  it('returns empty env object when no --env flags given', () => {
-    const o = parseTemplateOverrides(['--max-cost', '5']);
-    expect(o.env).toEqual({});
-  });
-});
-
-// ─────────────────────────────────────────────────────────────────────────────
-// 4. badgr serve template <name> — endpoint dispatch
-// ─────────────────────────────────────────────────────────────────────────────
-
-// Sequence for serveCommand:
-//   callApi('/serve', ...) → dep  (via callWithFallback → callApi)
-//   callApi('/deployments/<id>', ...) → { status: 'running' }  (pre-health check)
-//   callApi('/deployments/<id>', ...) → { status: 'running', endpoint_ready: true }  (waitForEndpoint poll)
-//
-// Readiness comes entirely from Badgr's own deployment status — the CLI never
-// fetches the pod/RunPod-proxy endpoint directly.
-
-function setupServe(depOverrides = {}, readyOverrides = {}) {
-  api.callApi
-    .mockResolvedValueOnce(makeServeDep(depOverrides))  // POST /serve
-    .mockResolvedValueOnce({ status: 'running' })        // pre-health dep status
-    .mockResolvedValueOnce({ status: 'running', endpoint_ready: true, ...readyOverrides }); // waitForEndpoint poll
-}
-
-describe('badgr serve template <name>', () => {
-  it('vllm: dispatches with correct image, max_cost_usd, and readiness on /v1/models', async () => {
-    setupServe({}, { health_path: '/v1/models' });
-    const p = serveCommand(config, ['template', 'vllm', '--max-cost', '5'], chalk);
-    await vi.advanceTimersByTimeAsync(5000);
-    await p;
-
-    const [route, opts] = api.callApi.mock.calls[0];
-    expect(route).toBe('/serve');
-    expect(opts.body.image).toBe('vllm/vllm-openai:latest');
-    expect(opts.body.max_cost_usd).toBe(5);
-    expect(process.exitCode).toBeFalsy();
-  });
-
-  it('comfyui: dispatches with ComfyUI image and readiness on /system_stats', async () => {
-    setupServe({}, { health_path: '/system_stats' });
-    const p = serveCommand(config, ['template', 'comfyui', '--max-cost', '3'], chalk);
-    await vi.advanceTimersByTimeAsync(5000);
-    await p;
-
-    const [, opts] = api.callApi.mock.calls[0];
-    expect(opts.body.image).toBe('yanwk/comfyui-boot:cu126-megapak');
-    expect(process.exitCode).toBeFalsy();
-  });
-
-  it('sglang: dispatches with A100 GPU', async () => {
-    setupServe({ gpu_type: 'A100' });
-    const p = serveCommand(config, ['template', 'sglang', '--max-cost', '10'], chalk);
-    await vi.advanceTimersByTimeAsync(5000);
-    await p;
-
-    const [, opts] = api.callApi.mock.calls[0];
-    expect(opts.body.image).toBe('lmsysorg/sglang:latest');
-    expect(opts.body.gpu).toBe('A100');
-    expect(process.exitCode).toBeFalsy();
-  });
-
-  it('tgi: dispatches with TGI image', async () => {
-    setupServe();
-    const p = serveCommand(config, ['template', 'tgi', '--max-cost', '5'], chalk);
-    await vi.advanceTimersByTimeAsync(5000);
-    await p;
-
-    const [, opts] = api.callApi.mock.calls[0];
-    expect(opts.body.image).toBe('ghcr.io/huggingface/text-generation-inference:latest');
-    expect(process.exitCode).toBeFalsy();
-  });
-
-  it('llama-cpp: dispatches with llama-cpp image', async () => {
-    setupServe();
-    const p = serveCommand(config, ['template', 'llama-cpp', '--max-cost', '3'], chalk);
-    await vi.advanceTimersByTimeAsync(5000);
-    await p;
-
-    const [, opts] = api.callApi.mock.calls[0];
-    expect(opts.body.image).toBe('michaelmanleyx/llama-cpp:server-cuda');
-    expect(process.exitCode).toBeFalsy();
-  });
-
-  it('invokeai: dispatches with InvokeAI image', async () => {
-    setupServe();
-    const p = serveCommand(config, ['template', 'invokeai', '--max-cost', '5'], chalk);
-    await vi.advanceTimersByTimeAsync(5000);
-    await p;
-
-    const [, opts] = api.callApi.mock.calls[0];
-    expect(opts.body.image).toBe('ghcr.io/invoke-ai/invokeai:latest');
-    expect(process.exitCode).toBeFalsy();
-  });
-
-  it('kohya-ss: dispatches with Kohya SS image', async () => {
-    setupServe();
-    const p = serveCommand(config, ['template', 'kohya-ss', '--max-cost', '5'], chalk);
-    await vi.advanceTimersByTimeAsync(5000);
-    await p;
-
-    const [, opts] = api.callApi.mock.calls[0];
-    expect(opts.body.image).toBe('bmaltais/kohya-ss-gui:latest');
-    expect(process.exitCode).toBeFalsy();
-  });
-
-  it('text-gen-webui: dispatches with text-gen-webui image', async () => {
-    setupServe();
-    const p = serveCommand(config, ['template', 'text-gen-webui', '--max-cost', '5'], chalk);
-    await vi.advanceTimersByTimeAsync(5000);
-    await p;
-
-    const [, opts] = api.callApi.mock.calls[0];
-    expect(opts.body.image).toBe('atinoda/text-generation-webui:default-nightly');
-    expect(process.exitCode).toBeFalsy();
-  });
-
-  it('user --gpu override is forwarded', async () => {
-    setupServe({ gpu_type: 'H100' });
-    const p = serveCommand(config, ['template', 'vllm', '--gpu', 'H100', '--max-cost', '20'], chalk);
-    await vi.advanceTimersByTimeAsync(5000);
-    await p;
-
-    const [, opts] = api.callApi.mock.calls[0];
-    expect(opts.body.gpu).toBe('H100');
-    expect(process.exitCode).toBeFalsy();
-  });
-
-  it('user --env override replaces template default', async () => {
-    setupServe();
-    const p = serveCommand(
-      config,
-      ['template', 'vllm', '--env', 'MODEL=mistralai/Mistral-7B-v0.1', '--max-cost', '5'],
-      chalk,
-    );
-    await vi.advanceTimersByTimeAsync(5000);
-    await p;
-
-    const [, opts] = api.callApi.mock.calls[0];
-    expect(opts.body.env.MODEL).toBe('mistralai/Mistral-7B-v0.1');
-    expect(process.exitCode).toBeFalsy();
-  });
-});
-
-// ─────────────────────────────────────────────────────────────────────────────
-// 5. badgr run template <name> — job dispatch
-// ─────────────────────────────────────────────────────────────────────────────
-
-// Sequence for runCommand:
-//   callApi('/run', ...) → dep  (via callWithFallback → callApi)
-//   callApi('/deployments/<id>', ...) → { status: 'completed', exit_code: 0 }
-//   callApi('/deployments/<id>/logs', ...) → { logs: [] }
-
-function setupRun(depOverrides = {}) {
-  api.callApi
-    .mockResolvedValueOnce(makeRunDep(depOverrides))                // POST /run
-    .mockResolvedValueOnce({ status: 'completed', exit_code: 0 })  // status poll
-    .mockResolvedValueOnce({ logs: [] });                           // logs
-}
-
-describe('badgr run template <name>', () => {
-  it('axolotl: dispatches with A100 image, max_cost, max_runtime', async () => {
-    setupRun();
-    const p = runCommand(config, ['template', 'axolotl', '--max-cost', '20', '--max-runtime', '120'], chalk);
-    await vi.advanceTimersByTimeAsync(10000);
-    await p;
-
-    const [route, opts] = api.callApi.mock.calls[0];
-    expect(route).toBe('/run');
-    expect(opts.body.image).toBe('axolotlai/axolotl-cloud-uv:main-latest');
-    expect(opts.body.gpu).toBe('A100');
-    expect(opts.body.max_cost_usd).toBe(20);
-    expect(opts.body.max_runtime_seconds).toBe(7200); // 120 min * 60
-    expect(process.exitCode).toBeFalsy();
-  });
-
-  it('unsloth: dispatches with unsloth image and RTX_4090', async () => {
-    setupRun({ gpu_type: 'RTX_4090' });
-    const p = runCommand(config, ['template', 'unsloth', '--max-cost', '10', '--max-runtime', '90'], chalk);
-    await vi.advanceTimersByTimeAsync(10000);
-    await p;
-
-    const [, opts] = api.callApi.mock.calls[0];
-    expect(opts.body.image).toBe('unslothai/unsloth:latest');
-    expect(opts.body.gpu).toBe('RTX_4090');
-    expect(process.exitCode).toBeFalsy();
-  });
-});
-
-// ─────────────────────────────────────────────────────────────────────────────
-// 6. Routing guards
-// ─────────────────────────────────────────────────────────────────────────────
-
-describe('routing guards', () => {
-  it('badgr run template <endpoint> → error with serve suggestion', async () => {
-    await runCommand(config, ['template', 'vllm', '--max-cost', '5'], chalk);
-    expect(process.exitCode).toBe(1);
-    expect(console.error).toHaveBeenCalledWith(
-      expect.stringContaining('badgr serve template vllm'),
-    );
-  });
-
-  it('badgr serve template <job> → error with run suggestion', async () => {
-    await serveCommand(config, ['template', 'axolotl', '--max-cost', '20'], chalk);
-    expect(process.exitCode).toBe(1);
-    expect(console.error).toHaveBeenCalledWith(
-      expect.stringContaining('badgr run template axolotl'),
-    );
-  });
-
-  it('badgr serve template <unknown> → error with list suggestion', async () => {
-    await serveCommand(config, ['template', 'doesnotexist'], chalk);
-    expect(process.exitCode).toBe(1);
-    expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Unknown template'));
-    expect(console.error).toHaveBeenCalledWith(expect.stringContaining('badgr template list'));
-  });
-
-  it('badgr run template <unknown> → error with list suggestion', async () => {
-    await runCommand(config, ['template', 'alsonotreal'], chalk);
-    expect(process.exitCode).toBe(1);
-    expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Unknown template'));
-  });
-
-  it('badgr serve template with no name → error', async () => {
-    await serveCommand(config, ['template'], chalk);
-    expect(process.exitCode).toBe(1);
-  });
-
-  it('badgr run template with no name → error', async () => {
-    await runCommand(config, ['template'], chalk);
-    expect(process.exitCode).toBe(1);
-  });
-});