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('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('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);

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

    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(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);
  });
});