badgr-cli 1.0.31 → 1.0.34

package/tests/run-lifecycle.test.js

@@ -0,0 +1,508 @@
+/**
+ * badgr run — end-to-end lifecycle tests (mocked API)
+ *
+ * Covers the launch-readiness matrix:
+ *   - Successful run on two GPU types (A100, L40S) via two providers (runpod, vastai)
+ *   - Job exits non-zero  → receipt failureType = customer_code
+ *   - Job exceeds max runtime  → teardown called, receipt reason = max-runtime
+ *   - Job exceeds max cost     → teardown called, receipt reason = max-cost
+ *   - Container fails before running (infrastructure)
+ *   - Heartbeat lost mid-run   → teardown called, receipt reason = heartbeat-lost
+ *   - Log streaming            → log lines appear in console output
+ *   - Tier-1 unavailable       → auto-expand to tier-2, receipt records tier
+ *   - No capacity at all       → CapacityError displayed, exitCode = 1
+ *   - Billing stops            → terminateDeployment called on every terminal path
+ *   - Receipt explains state   → status, exitCode, failureType, runtimeSeconds, finalCost
+ */
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { runCommand } from '../src/commands/run.js';
+
+// ── Module mocks ──────────────────────────────────────────────────────────────
+
+vi.mock('../src/api.js', () => ({
+  callApi:             vi.fn(),
+  terminateDeployment: vi.fn().mockResolvedValue({}),
+}));
+
+vi.mock('../src/store.js', () => ({
+  addDeployment:        vi.fn(),
+  addReceipt:           vi.fn(),
+  updateReceipt:        vi.fn(),
+  generateReceiptId:    vi.fn(() => 'rcpt-launch-001'),
+  generateDeploymentId: vi.fn(() => 'dep-launch-001'),
+  listDeployments:      vi.fn(() => []),
+  listReceipts:         vi.fn(() => []),
+  findDeployment:       vi.fn(() => null),
+  updateDeployment:     vi.fn(),
+  removeDeployment:     vi.fn(),
+}));
+
+// Imported after mock declarations so we get the mocked versions
+import * as api   from '../src/api.js';
+import * as store from '../src/store.js';
+
+// ── Helpers ───────────────────────────────────────────────────────────────────
+
+const config = { apiKey: 'sk-test', baseUrl: 'https://api.test/v1' };
+
+// Chalk no-op (the real chalk would add ANSI codes, which makes assertions noisy)
+const chalk = {
+  bold:   s => s,
+  dim:    s => s,
+  red:    s => s,
+  yellow: s => s,
+  green:  s => s,
+  cyan:   s => s,
+};
+
+function makeDep(overrides = {}) {
+  return {
+    deployment_id: 'dep-launch-001',
+    status:        'running',
+    gpu_type:      'A100',
+    gpu_count:     1,
+    cost_per_hour: 2.50,
+    provider:      'runpod',
+    receipt_id:    'rcpt-launch-001',
+    tier:          '1',
+    ...overrides,
+  };
+}
+
+// Standard happy-path sequence: POST /run → GET /deployments → GET /logs → (no DELETE needed for success)
+function setupSuccessfulRun(depOverrides = {}, exitCode = 0) {
+  api.callApi
+    .mockResolvedValueOnce(makeDep(depOverrides))                         // POST /run
+    .mockResolvedValueOnce({ status: 'completed', exit_code: exitCode }) // GET /deployments/…
+    .mockResolvedValueOnce({ logs: ['step 1', 'step 2'] });              // GET /deployments/…/logs
+  api.terminateDeployment.mockResolvedValue({});
+}
+
+// ── Setup / teardown ─────────────────────────────────────────────────────────
+
+// Suppress MaxListenersExceededWarning — each runCommand call adds a SIGINT listener
+// and we run many tests in sequence.
+process.setMaxListeners(50);
+
+beforeEach(() => {
+  // Explicitly list what to fake — vitest 4.x also fakes queueMicrotask and
+  // nextTick by default which deadlocks await-inside-while-loop patterns.
+  vi.useFakeTimers({
+    toFake: ['setTimeout', 'setInterval', 'clearTimeout', 'clearInterval', 'Date', 'performance', 'setImmediate', 'clearImmediate'],
+  });
+  process.exitCode = undefined;
+  vi.spyOn(console, 'log').mockImplementation(() => {});
+  vi.spyOn(console, 'error').mockImplementation(() => {});
+  vi.spyOn(process.stdout, 'write').mockImplementation(() => true);
+  // resetAllMocks clears both call history AND the mockResolvedValueOnce queue
+  // (clearAllMocks only clears history in vitest 4.x, leaving stale queued values
+  // that can corrupt subsequent tests).
+  vi.resetAllMocks();
+  // Restore default mock return values cleared by resetAllMocks
+  store.generateReceiptId.mockReturnValue('rcpt-launch-001');
+  store.generateDeploymentId.mockReturnValue('dep-launch-001');
+  store.listDeployments.mockReturnValue([]);
+  store.listReceipts.mockReturnValue([]);
+  store.findDeployment.mockReturnValue(null);
+  api.terminateDeployment.mockResolvedValue({});
+});
+
+afterEach(() => {
+  vi.useRealTimers();
+  vi.restoreAllMocks();
+  process.exitCode = undefined;
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 1. GPU / provider matrix
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('GPU / provider matrix', () => {
+  it('A100 via RunPod (tier-1) completes successfully', async () => {
+    setupSuccessfulRun({ gpu_type: 'A100', provider: 'runpod', tier: '1' });
+    const p = runCommand(config, ['python', 'train.py', '--gpu', 'A100'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    expect(api.callApi).toHaveBeenCalledWith('/run', expect.objectContaining({ method: 'POST' }));
+    expect(store.addReceipt).toHaveBeenCalledWith(expect.objectContaining({
+      gpu:           'A100',
+      providerRoute: 'runpod',
+      tier:          '1',
+    }));
+    expect(store.updateReceipt).toHaveBeenCalledWith('rcpt-launch-001', expect.objectContaining({
+      status:   'completed',
+      exitCode: 0,
+    }));
+    expect(process.exitCode).toBeFalsy();
+  });
+
+  it('L40S via Vast.ai (tier-1) completes successfully', async () => {
+    api.callApi
+      .mockResolvedValueOnce(makeDep({ gpu_type: 'L40S', provider: 'vastai', tier: '1', cost_per_hour: 1.80 }))
+      .mockResolvedValueOnce({ status: 'completed', exit_code: 0 })
+      .mockResolvedValueOnce({ logs: [] });
+    const p = runCommand(config, ['python', 'train.py', '--gpu', 'L40S'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    expect(store.addReceipt).toHaveBeenCalledWith(expect.objectContaining({
+      gpu:           'L40S',
+      providerRoute: 'vastai',
+    }));
+    expect(store.updateReceipt).toHaveBeenCalledWith('rcpt-launch-001', expect.objectContaining({
+      status:   'completed',
+      exitCode: 0,
+    }));
+    expect(process.exitCode).toBeFalsy();
+  });
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 2. Job exits non-zero
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('non-zero exit code', () => {
+  it('records customer_code failure and sets exitCode', async () => {
+    api.callApi
+      .mockResolvedValueOnce(makeDep())
+      .mockResolvedValueOnce({ status: 'failed', exit_code: 1 })
+      .mockResolvedValueOnce({ logs: ['Traceback (most recent call last)'] });
+    const p = runCommand(config, ['python', 'train.py'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    expect(store.updateReceipt).toHaveBeenCalledWith('rcpt-launch-001', expect.objectContaining({
+      exitCode:    1,
+      failureType: 'customer_code',
+    }));
+    expect(process.exitCode).toBe(1);
+  });
+
+  it('distinguishes customer_code from infrastructure when exit code is null', async () => {
+    api.callApi
+      .mockResolvedValueOnce(makeDep())
+      .mockResolvedValueOnce({ status: 'failed', exit_code: null })
+      .mockResolvedValueOnce({ logs: [] });
+    const p = runCommand(config, ['python', 'train.py'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    expect(store.updateReceipt).toHaveBeenCalledWith('rcpt-launch-001', expect.objectContaining({
+      failureType: 'infrastructure',
+    }));
+  });
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 3. Container fails to start (infrastructure failure before running)
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('container fails to start', () => {
+  it('reports infrastructure error when dep arrives already failed', async () => {
+    api.callApi
+      .mockResolvedValueOnce(makeDep({ status: 'failed' })) // POST /run returns already-failed
+      .mockResolvedValueOnce({ logs: [] });
+    const p = runCommand(config, ['python', 'train.py'], chalk);
+    await vi.advanceTimersByTimeAsync(1000);
+    await p;
+
+    // runCommand exits with process.exitCode = 1 when dep.status === 'failed' after provision
+    expect(process.exitCode).toBe(1);
+  });
+
+  it('reports infrastructure error when waitForRunning resolves to failed', async () => {
+    api.callApi
+      .mockResolvedValueOnce(makeDep({ status: 'provisioning' })) // POST /run → provisioning
+      .mockResolvedValueOnce({ status: 'failed' });               // poll → failed
+    const p = runCommand(config, ['python', 'train.py'], chalk);
+    // advance past POLL_MS (3000ms) in waitForRunning
+    await vi.advanceTimersByTimeAsync(4000);
+    await p;
+
+    expect(process.exitCode).toBe(1);
+  });
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 4. Max-runtime cap
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('max-runtime cap', () => {
+  it('stops job and records max-runtime when runtime exceeded', async () => {
+    // max-runtime 0.05 min = 3 seconds; ratePerHour = 2.50
+    api.callApi
+      .mockResolvedValueOnce(makeDep({ cost_per_hour: 2.50 }))  // POST /run
+      .mockResolvedValueOnce({ status: 'running', exit_code: null }) // poll 1
+      .mockResolvedValueOnce({ logs: [] });                          // logs 1
+    api.terminateDeployment.mockResolvedValue({});
+
+    const p = runCommand(config, ['python', 'train.py', '--max-runtime', '0.05'], chalk);
+    // attachToJob POLL_MS = 4000, maxRuntime = 0.05 min = 3000ms → cap fires after 4000ms (first elapsedMs check)
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    expect(api.terminateDeployment).toHaveBeenCalledWith(config, 'dep-launch-001');
+    expect(store.updateReceipt).toHaveBeenCalledWith('rcpt-launch-001', expect.objectContaining({
+      status:         'max-runtime',
+      teardownStatus: 'terminated',
+    }));
+    expect(process.exitCode).toBe(1);
+  });
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 5. Max-cost cap
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('max-cost cap', () => {
+  it('stops job and records max-cost when budget exceeded', async () => {
+    // ratePerHour = 3600 ($1/s), maxCost = 0.001 → cap fires after 3.6ms, but poll fires at 4000ms
+    api.callApi
+      .mockResolvedValueOnce(makeDep({ cost_per_hour: 3600 })) // POST /run
+      .mockResolvedValueOnce({ status: 'running' })            // poll 1 (never reached — cap fires first)
+      .mockResolvedValueOnce({ logs: [] });
+    api.terminateDeployment.mockResolvedValue({});
+
+    const p = runCommand(config, ['python', 'train.py', '--max-cost', '0.001'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    expect(api.terminateDeployment).toHaveBeenCalledWith(config, 'dep-launch-001');
+    expect(store.updateReceipt).toHaveBeenCalledWith('rcpt-launch-001', expect.objectContaining({
+      status:         'max-cost',
+      teardownStatus: 'terminated',
+    }));
+    expect(process.exitCode).toBe(1);
+  });
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 6. Heartbeat lost
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('heartbeat lost', () => {
+  it('stops job and records heartbeat-lost after HEARTBEAT_KILL_POLLS consecutive errors', async () => {
+    // HEARTBEAT_KILL_POLLS = 15; POLL_MS = 4000ms
+    // First dep poll must succeed and return 'running' (to set lastStatus = 'running').
+    // After that, 15 consecutive failures trigger the kill.
+    // Total timer advances: (1 success + 15 failures) × 4000ms = 64000ms
+    api.callApi
+      .mockResolvedValueOnce(makeDep({ status: 'running' })) // POST /run
+      .mockResolvedValueOnce({ status: 'running' })          // first dep poll → sets lastStatus = 'running'
+      .mockResolvedValueOnce({ logs: [] })                   // first logs poll
+      .mockRejectedValue(new Error('ETIMEDOUT'));             // all subsequent dep/logs polls fail
+    api.terminateDeployment.mockResolvedValue({});
+
+    const p = runCommand(config, ['python', 'train.py'], chalk);
+    // Advance past 1 successful poll (4000ms) + 15 failed polls (60000ms) = 64000ms
+    await vi.advanceTimersByTimeAsync(66000);
+    await p;
+
+    expect(api.terminateDeployment).toHaveBeenCalledWith(config, 'dep-launch-001');
+
+    // teardown() writes first: billing stopped, status = 'heartbeat-lost'
+    expect(store.updateReceipt).toHaveBeenCalledWith('rcpt-launch-001', expect.objectContaining({
+      status:         'heartbeat-lost',
+      teardownStatus: 'terminated',
+    }));
+    // runCommand writes second: final summary with failureType
+    expect(store.updateReceipt).toHaveBeenCalledWith('rcpt-launch-001', expect.objectContaining({
+      failureType: 'infrastructure',
+    }));
+    expect(process.exitCode).toBeTruthy();
+  }, 30000 /* 30s real timeout — fake timers advance instantly but need sufficient budget */);
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 7. Log streaming
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('log streaming', () => {
+  it('prints user-visible log lines and suppresses meta-lines', async () => {
+    const logLines = [];
+    console.log.mockImplementation((...args) => logLines.push(args.join(' ')));
+
+    api.callApi
+      .mockResolvedValueOnce(makeDep())
+      .mockResolvedValueOnce({ status: 'completed', exit_code: 0 })
+      .mockResolvedValueOnce({ logs: [
+        'Epoch 1/3: loss=0.42',                         // should appear
+        '[dep-abc] status=running',                     // meta → should be filtered
+        '[dep-abc] gpu_util=72%',                       // meta → should be filtered
+        'Epoch 2/3: loss=0.31',                         // should appear
+      ]});
+
+    const p = runCommand(config, ['python', 'train.py'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    const joined = logLines.join('\n');
+    expect(joined).toContain('Epoch 1/3: loss=0.42');
+    expect(joined).toContain('Epoch 2/3: loss=0.31');
+    // Meta-lines are filtered by the LOG_META_RE in run.js
+    expect(joined).not.toContain('[dep-abc] status=running');
+  });
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 8. Routing: tier-1 unavailable → auto-expand to tier-2
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('routing fallback', () => {
+  it('expands to tier-2 when tier-1 returns NO_CAPACITY_MATCH', async () => {
+    const capacityErr = Object.assign(new Error('no capacity'), {
+      errorData: { code: 'NO_CAPACITY_MATCH' },
+      httpStatus: 503,
+    });
+    api.callApi
+      .mockRejectedValueOnce(capacityErr)                                     // tier-1 attempt
+      .mockResolvedValueOnce(makeDep({ provider: 'modal', tier: '2' }))      // tier-2 succeeds
+      .mockResolvedValueOnce({ status: 'completed', exit_code: 0 })
+      .mockResolvedValueOnce({ logs: [] });
+
+    const p = runCommand(config, ['python', 'train.py'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    // Should have been called twice: once for tier-1, once for tier-2
+    const postCalls = api.callApi.mock.calls.filter(c => c[1]?.method === 'POST');
+    expect(postCalls.length).toBeGreaterThanOrEqual(2);
+
+    expect(store.addReceipt).toHaveBeenCalledWith(expect.objectContaining({
+      providerRoute: 'modal',
+      tier:          '2',
+    }));
+    expect(process.exitCode).toBeFalsy();
+  });
+
+  it('reports CapacityError when both tiers fail', async () => {
+    const capacityErr = Object.assign(new Error('no capacity'), {
+      errorData: { code: 'NO_CAPACITY_MATCH' },
+      httpStatus: 503,
+    });
+    api.callApi.mockRejectedValue(capacityErr);
+
+    const logs = [];
+    console.error.mockImplementation(msg => logs.push(msg));
+
+    await runCommand(config, ['python', 'train.py'], chalk);
+
+    expect(process.exitCode).toBe(1);
+    const combined = logs.join('\n');
+    expect(combined).toMatch(/No suitable GPU capacity|capacity/i);
+  });
+
+  it('does not fall back to tier-2 when --no-fallback is set', async () => {
+    const capacityErr = Object.assign(new Error('no capacity'), {
+      errorData: { code: 'NO_CAPACITY_MATCH' },
+      httpStatus: 503,
+    });
+    api.callApi.mockRejectedValueOnce(capacityErr);
+
+    await runCommand(config, ['python', 'train.py', '--no-fallback'], chalk);
+
+    // Only one POST /run call was made (no tier-2 expansion)
+    const postCalls = api.callApi.mock.calls.filter(c => c[1]?.method === 'POST');
+    expect(postCalls.length).toBe(1);
+    expect(process.exitCode).toBe(1);
+  });
+
+  it('PROVISIONING_FAILED error is surfaced clearly without showing internal debug info by default', async () => {
+    const provErr = Object.assign(new Error('provisioning failed'), {
+      errorData: { code: 'PROVISIONING_FAILED', debug_error: 'CUDA driver mismatch' },
+      httpStatus: 503,
+    });
+    api.callApi.mockRejectedValue(provErr);
+
+    const errLines = [];
+    console.error.mockImplementation(msg => errLines.push(msg));
+
+    await runCommand(config, ['python', 'train.py'], chalk);
+
+    expect(process.exitCode).toBe(1);
+    const combined = errLines.join('\n');
+    expect(combined).not.toContain('CUDA driver mismatch'); // debug info hidden by default
+    expect(combined).toMatch(/try again|failed to start/i);
+  });
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 9. Payment required
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('payment required', () => {
+  it('shows billing message and rerun hint, exits 1', async () => {
+    const payErr = Object.assign(new Error('Payment required'), {
+      isPaymentRequired: true,
+      httpStatus: 402,
+    });
+    api.callApi.mockRejectedValueOnce(payErr);
+
+    const errLines = [];
+    console.error.mockImplementation(msg => errLines.push(msg));
+
+    await runCommand(config, ['python', 'train.py', '--gpu', 'A100'], chalk);
+
+    expect(process.exitCode).toBe(1);
+    const combined = errLines.join('\n');
+    expect(combined).toContain('Payment required');
+  });
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 10. Billing stops — terminateDeployment called on every terminal path
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('billing lifecycle', () => {
+  it('terminateDeployment is called after successful completion', async () => {
+    setupSuccessfulRun();
+    const p = runCommand(config, ['python', 'train.py'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+    expect(api.terminateDeployment).toHaveBeenCalledWith(config, 'dep-launch-001');
+  });
+
+  it('receipt records runtime and finalCost on completion', async () => {
+    setupSuccessfulRun({ cost_per_hour: 1.80 });
+    const p = runCommand(config, ['python', 'train.py'], chalk);
+    await vi.advanceTimersByTimeAsync(5000);
+    await p;
+
+    const updateCall = store.updateReceipt.mock.calls.find(
+      c => c[1]?.runtimeSeconds !== undefined,
+    );
+    expect(updateCall).toBeTruthy();
+    expect(typeof updateCall[1].runtimeSeconds).toBe('number');
+    expect(typeof updateCall[1].finalCost).toBe('number');
+  });
+});
+
+// ─────────────────────────────────────────────────────────────────────────────
+// 11. Input validation before provisioning
+// ─────────────────────────────────────────────────────────────────────────────
+
+describe('input validation', () => {
+  it('rejects --count 0 before any API call', async () => {
+    await runCommand(config, ['python', 'train.py', '--count', '0'], chalk);
+    expect(api.callApi).not.toHaveBeenCalled();
+    expect(process.exitCode).toBe(1);
+  });
+
+  it('rejects --max-cost 0 before any API call', async () => {
+    await runCommand(config, ['python', 'train.py', '--max-cost', '0'], chalk);
+    expect(api.callApi).not.toHaveBeenCalled();
+    expect(process.exitCode).toBe(1);
+  });
+
+  it('rejects unknown --region before any API call', async () => {
+    await runCommand(config, ['python', 'train.py', '--region', 'MARS'], chalk);
+    expect(api.callApi).not.toHaveBeenCalled();
+    expect(process.exitCode).toBe(1);
+  });
+
+  it('rejects empty command with no --image', async () => {
+    await runCommand(config, [], chalk);
+    expect(api.callApi).not.toHaveBeenCalled();
+  });
+});