flynn/src/companion/heartbeatLoop.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { CompanionHeartbeatLoop } from './heartbeatLoop.js';
import type { NodeStatusSetResult } from './runtimeClient.js';

function buildStatusResult(): NodeStatusSetResult {
  return {
    updated: true,
    node: { id: 'node-1', role: 'companion' },
    status: {
      platform: 'macos',
      powerSource: 'unknown',
      reportedAt: Date.now(),
    },
  };
}

describe('CompanionHeartbeatLoop', () => {
  beforeEach(() => {
    vi.useFakeTimers();
  });

  afterEach(() => {
    vi.useRealTimers();
  });

  it('runs immediately by default and continues on interval', async () => {
    const publishHeartbeat = vi.fn(async () => buildStatusResult());
    const loop = new CompanionHeartbeatLoop({ publishHeartbeat }, { intervalMs: 1000 });

    loop.start();
    await Promise.resolve();
    expect(publishHeartbeat).toHaveBeenCalledTimes(1);

    await vi.advanceTimersByTimeAsync(1000);
    expect(publishHeartbeat).toHaveBeenCalledTimes(2);

    loop.stop();
  });

  it('validates constructor options', () => {
    const publisher = { publishHeartbeat: vi.fn(async () => buildStatusResult()) };
    expect(() => new CompanionHeartbeatLoop(publisher, { intervalMs: 0 })).toThrow(
      'intervalMs must be a positive number',
    );
    expect(() => new CompanionHeartbeatLoop(publisher, { jitterRatio: -0.1 })).toThrow(
      'jitterRatio must be between 0 and 1',
    );
    expect(() => new CompanionHeartbeatLoop(publisher, { jitterRatio: 1.1 })).toThrow(
      'jitterRatio must be between 0 and 1',
    );
    expect(() => new CompanionHeartbeatLoop(publisher, { maxConsecutiveFailures: 0 })).toThrow(
      'maxConsecutiveFailures must be >= 1 when specified',
    );
  });

  it('supports delayed first run when runImmediately=false', async () => {
    const publishHeartbeat = vi.fn(async () => buildStatusResult());
    const loop = new CompanionHeartbeatLoop({ publishHeartbeat }, { intervalMs: 500 });

    loop.start(false);
    expect(publishHeartbeat).toHaveBeenCalledTimes(0);

    await vi.advanceTimersByTimeAsync(500);
    expect(publishHeartbeat).toHaveBeenCalledTimes(1);

    loop.stop();
  });

  it('applies jitterRatio to scheduling delay bounds', async () => {
    const publishHeartbeat = vi.fn(async () => buildStatusResult());
    const loop = new CompanionHeartbeatLoop(
      { publishHeartbeat },
      { intervalMs: 1000, jitterRatio: 0.5, randomFn: () => 1 },
    );

    loop.start(false);
    await vi.advanceTimersByTimeAsync(1499);
    expect(publishHeartbeat).toHaveBeenCalledTimes(0);
    await vi.advanceTimersByTimeAsync(1);
    expect(publishHeartbeat).toHaveBeenCalledTimes(1);
    loop.stop();
  });

  it('normalizes invalid randomFn samples for jitter scheduling', async () => {
    const publishHeartbeatHigh = vi.fn(async () => buildStatusResult());
    const loopHigh = new CompanionHeartbeatLoop(
      { publishHeartbeat: publishHeartbeatHigh },
      { intervalMs: 1000, jitterRatio: 0.5, randomFn: () => 2 },
    );
    loopHigh.start(false);
    await vi.advanceTimersByTimeAsync(1499);
    expect(publishHeartbeatHigh).toHaveBeenCalledTimes(0);
    await vi.advanceTimersByTimeAsync(1);
    expect(publishHeartbeatHigh).toHaveBeenCalledTimes(1);
    loopHigh.stop();

    const publishHeartbeatNaN = vi.fn(async () => buildStatusResult());
    const loopNaN = new CompanionHeartbeatLoop(
      { publishHeartbeat: publishHeartbeatNaN },
      { intervalMs: 1000, jitterRatio: 0.5, randomFn: () => Number.NaN },
    );
    loopNaN.start(false);
    await vi.advanceTimersByTimeAsync(999);
    expect(publishHeartbeatNaN).toHaveBeenCalledTimes(0);
    await vi.advanceTimersByTimeAsync(1);
    expect(publishHeartbeatNaN).toHaveBeenCalledTimes(1);
    loopNaN.stop();
  });

  it('calls onSuccess with heartbeat result payload', async () => {
    const result = buildStatusResult();
    const publishHeartbeat = vi.fn(async () => result);
    const onSuccess = vi.fn();
    const loop = new CompanionHeartbeatLoop(
      { publishHeartbeat },
      { intervalMs: 200, onSuccess },
    );

    loop.start();
    await Promise.resolve();
    expect(onSuccess).toHaveBeenCalledWith(result);
    loop.stop();
  });

  it('tracks successCount and lastSuccessAt in loop state', async () => {
    const publishHeartbeat = vi.fn(async () => buildStatusResult());
    const loop = new CompanionHeartbeatLoop({ publishHeartbeat }, { intervalMs: 300 });

    expect(loop.successCount).toBe(0);
    expect(loop.lastSuccessAt).toBeNull();

    loop.start();
    await Promise.resolve();

    expect(loop.successCount).toBe(1);
    expect(loop.lastSuccessAt).not.toBeNull();
    expect(loop.getState().successCount).toBe(1);
    expect(loop.getState().lastSuccessAt).toBe(loop.lastSuccessAt);

    loop.stop();
  });

  it('passes buildHeartbeat payload into publishHeartbeat', async () => {
    const publishHeartbeat = vi.fn(async () => buildStatusResult());
    const buildHeartbeat = vi.fn(() => ({ statusText: 'loop', powerSource: 'ac' as const }));
    const loop = new CompanionHeartbeatLoop(
      { publishHeartbeat },
      { intervalMs: 500, buildHeartbeat },
    );

    loop.start();
    await Promise.resolve();

    expect(buildHeartbeat).toHaveBeenCalledTimes(1);
    expect(publishHeartbeat).toHaveBeenCalledWith({ statusText: 'loop', powerSource: 'ac' });

    loop.stop();
  });

  it('reports errors through onError and keeps scheduling', async () => {
    const publishHeartbeat = vi
      .fn<() => Promise<NodeStatusSetResult>>()
      .mockRejectedValueOnce(new Error('boom'))
      .mockResolvedValue(buildStatusResult());
    const onError = vi.fn();
    const loop = new CompanionHeartbeatLoop({ publishHeartbeat }, { intervalMs: 400, onError });

    loop.start();
    await Promise.resolve();
    expect(onError).toHaveBeenCalledTimes(1);
    expect(loop.failureCount).toBe(1);
    expect(loop.lastFailure?.message).toBe('boom');

    await vi.advanceTimersByTimeAsync(400);
    expect(publishHeartbeat).toHaveBeenCalledTimes(2);
    expect(loop.failureCount).toBe(0);
    expect(loop.lastFailure).toBeNull();

    loop.stop();
  });

  it('is idempotent for repeated start and stop calls', async () => {
    const publishHeartbeat = vi.fn(async () => buildStatusResult());
    const loop = new CompanionHeartbeatLoop({ publishHeartbeat }, { intervalMs: 250 });

    loop.start();
    loop.start();
    await Promise.resolve();
    expect(publishHeartbeat).toHaveBeenCalledTimes(1);

    loop.stop();
    loop.stop();
    await vi.advanceTimersByTimeAsync(1000);
    expect(publishHeartbeat).toHaveBeenCalledTimes(1);
    expect(loop.running).toBe(false);
  });

  it('stops when maxConsecutiveFailures threshold is reached', async () => {
    const publishHeartbeat = vi
      .fn<() => Promise<NodeStatusSetResult>>()
      .mockRejectedValue(new Error('persistent-failure'));
    const onError = vi.fn();
    const onFailureLimitReached = vi.fn();
    const loop = new CompanionHeartbeatLoop(
      { publishHeartbeat },
      {
        intervalMs: 300,
        onError,
        maxConsecutiveFailures: 2,
        onFailureLimitReached,
      },
    );

    loop.start();
    await Promise.resolve();
    expect(loop.running).toBe(true);

    await vi.advanceTimersByTimeAsync(300);
    expect(loop.running).toBe(false);
    expect(loop.failureCount).toBe(2);
    expect(loop.lastFailure?.message).toBe('persistent-failure');
    expect(onError).toHaveBeenCalledTimes(2);
    expect(onFailureLimitReached).toHaveBeenCalledTimes(1);
    expect(onFailureLimitReached).toHaveBeenCalledWith(expect.any(Error), 2);

    await vi.advanceTimersByTimeAsync(1000);
    expect(publishHeartbeat).toHaveBeenCalledTimes(2);
  });

  it('tickNow sends heartbeat even when loop is not started', async () => {
    const publishHeartbeat = vi.fn(async () => buildStatusResult());
    const loop = new CompanionHeartbeatLoop({ publishHeartbeat }, { intervalMs: 1000 });

    expect(loop.running).toBe(false);
    await loop.tickNow();
    expect(publishHeartbeat).toHaveBeenCalledTimes(1);
    expect(loop.running).toBe(false);
    expect(loop.getState()).toEqual({
      running: false,
      successCount: 1,
      lastSuccessAt: expect.any(Number),
      failureCount: 0,
      lastFailure: null,
    });
  });

  it('tickNow while running does not accumulate duplicate scheduled timers', async () => {
    const publishHeartbeat = vi.fn(async () => buildStatusResult());
    const loop = new CompanionHeartbeatLoop({ publishHeartbeat }, { intervalMs: 1000 });

    loop.start(false);
    expect(vi.getTimerCount()).toBe(1);

    await loop.tickNow();
    expect(vi.getTimerCount()).toBe(1);
    expect(publishHeartbeat).toHaveBeenCalledTimes(1);

    await vi.advanceTimersByTimeAsync(1000);
    expect(publishHeartbeat).toHaveBeenCalledTimes(2);

    loop.stop();
  });

  it('resets success and failure state when restarted', async () => {
    const publishHeartbeat = vi
      .fn<() => Promise<NodeStatusSetResult>>()
      .mockResolvedValueOnce(buildStatusResult())
      .mockRejectedValueOnce(new Error('transient'))
      .mockResolvedValue(buildStatusResult());
    const loop = new CompanionHeartbeatLoop({ publishHeartbeat }, { intervalMs: 500 });

    loop.start();
    await Promise.resolve();
    expect(loop.successCount).toBe(1);
    expect(loop.failureCount).toBe(0);

    await vi.advanceTimersByTimeAsync(500);
    expect(loop.failureCount).toBe(1);
    expect(loop.lastFailure?.message).toBe('transient');

    loop.stop();
    loop.start(false);

    expect(loop.successCount).toBe(0);
    expect(loop.lastSuccessAt).toBeNull();
    expect(loop.failureCount).toBe(0);
    expect(loop.lastFailure).toBeNull();

    await vi.advanceTimersByTimeAsync(500);
    expect(loop.successCount).toBe(1);
    expect(loop.failureCount).toBe(0);

    loop.stop();
  });
});