flynn/src/cli/companion.test.ts

import { Command } from 'commander';
import { beforeEach, describe, expect, it, vi } from 'vitest';

const {
  mockLoadConfigSafe,
  mockGetConfigPath,
  mockRuntimeCtorArgs,
  mockRuntimeInstances,
} = vi.hoisted(() => {
  const runtimeCtorArgs: Array<{ url: string; token?: string; autoReconnect?: boolean }> = [];
  const runtimeInstances: Array<{
    connect: ReturnType<typeof vi.fn>;
    registerNode: ReturnType<typeof vi.fn>;
    setNodeStatus: ReturnType<typeof vi.fn>;
    sendAgentMessage: ReturnType<typeof vi.fn>;
    subscribeAgentStream: ReturnType<typeof vi.fn>;
    subscribeAgentTyping: ReturnType<typeof vi.fn>;
    subscribeConnectionEvents: ReturnType<typeof vi.fn>;
    disconnect: ReturnType<typeof vi.fn>;
  }> = [];

  const loadConfigSafe = vi.fn(() => ({
    config: {
      server: {
        port: 18888,
        token: 'config-token',
      },
    },
  }));

  const getConfigPath = vi.fn(() => '/tmp/flynn-config.yaml');

  return {
    mockLoadConfigSafe: loadConfigSafe,
    mockGetConfigPath: getConfigPath,
    mockRuntimeCtorArgs: runtimeCtorArgs,
    mockRuntimeInstances: runtimeInstances,
  };
});

vi.mock('./shared.js', () => ({
  loadConfigSafe: mockLoadConfigSafe,
  getConfigPath: mockGetConfigPath,
}));

vi.mock('../companion/index.js', () => ({
  CompanionRuntimeClient: class {
    private connectionHandlers: Array<(event: { status: string }) => void> = [];
    connect = vi.fn(async () => {
      for (const handler of this.connectionHandlers) {
        handler({ status: 'connected' });
      }
      return undefined;
    });
    registerNode = vi.fn(async ({ nodeId, role, capabilities }: { nodeId: string; role: string; capabilities: string[] }) => ({
      registered: true,
      node: { id: nodeId, role },
      protocol: { serverVersion: 1, clientVersion: 1, negotiatedVersion: 1 },
      capabilities: { declared: capabilities, enabled: capabilities },
    }));
    setNodeStatus = vi.fn(async () => ({ updated: true, node: { id: 'n', role: 'companion' } }));
    sendAgentMessage = vi.fn(async () => ({ content: 'handoff response' }));
    subscribeAgentStream = vi.fn(() => () => undefined);
    subscribeAgentTyping = vi.fn(() => () => undefined);
    subscribeConnectionEvents = vi.fn((handler: (event: { status: string }) => void) => {
      this.connectionHandlers.push(handler);
      return () => undefined;
    });
    disconnect = vi.fn(() => undefined);

    constructor(opts: { url: string; token?: string; autoReconnect?: boolean }) {
      mockRuntimeCtorArgs.push(opts);
      mockRuntimeInstances.push(this);
    }
  },
}));

describe('companion command', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mockRuntimeCtorArgs.length = 0;
    mockRuntimeInstances.length = 0;
    mockLoadConfigSafe.mockReturnValue({
      config: {
        server: {
          port: 18888,
          token: 'config-token',
        },
      },
    });
    mockGetConfigPath.mockReturnValue('/tmp/flynn-config.yaml');
    process.exitCode = undefined;
  });

  it('uses config-derived gateway url/token by default', async () => {
    const logSpy = vi.spyOn(console, 'log').mockImplementation(() => undefined);
    const errSpy = vi.spyOn(console, 'error').mockImplementation(() => undefined);
    const program = new Command();
    const { registerCompanionCommand } = await import('./companion.js');
    registerCompanionCommand(program);

    await program.parseAsync(['node', 'test', 'companion', '--once']);

    expect(mockGetConfigPath).toHaveBeenCalledOnce();
    expect(mockRuntimeCtorArgs).toEqual([{ url: 'ws://127.0.0.1:18888', token: 'config-token', autoReconnect: false }]);
    expect(mockRuntimeInstances[0]?.connect).toHaveBeenCalledOnce();
    expect(mockRuntimeInstances[0]?.registerNode).toHaveBeenCalledOnce();
    expect(mockRuntimeInstances[0]?.setNodeStatus).toHaveBeenCalledOnce();
    expect(mockRuntimeInstances[0]?.disconnect).toHaveBeenCalled();
    expect(errSpy).not.toHaveBeenCalled();
    expect(process.exitCode).toBeUndefined();

    logSpy.mockRestore();
    errSpy.mockRestore();
  });

  it('prefers explicit url/token and capability overrides', async () => {
    const logSpy = vi.spyOn(console, 'log').mockImplementation(() => undefined);
    const errSpy = vi.spyOn(console, 'error').mockImplementation(() => undefined);
    const program = new Command();
    const { registerCompanionCommand } = await import('./companion.js');
    registerCompanionCommand(program);

    await program.parseAsync([
      'node',
      'test',
      'companion',
      '--once',
      '--url',
      'ws://10.0.0.5:19000',
      '--token',
      'override-token',
      '--node-id',
      'test-node',
      '--capability',
      'ui.canvas',
      'node.push.register',
    ]);

    expect(mockRuntimeCtorArgs).toEqual([{ url: 'ws://10.0.0.5:19000', token: 'override-token', autoReconnect: false }]);
    expect(mockRuntimeInstances[0]?.registerNode).toHaveBeenCalledWith(expect.objectContaining({
      nodeId: 'test-node',
      capabilities: ['ui.canvas', 'node.push.register'],
    }));
    expect(errSpy).not.toHaveBeenCalled();

    logSpy.mockRestore();
    errSpy.mockRestore();
  });

  it('executes optional message handoff after registration', async () => {
    const logSpy = vi.spyOn(console, 'log').mockImplementation(() => undefined);
    const errSpy = vi.spyOn(console, 'error').mockImplementation(() => undefined);
    const program = new Command();
    const { registerCompanionCommand } = await import('./companion.js');
    registerCompanionCommand(program);

    await program.parseAsync([
      'node',
      'test',
      'companion',
      '--once',
      '--handoff',
      'status update?',
      '--handoff-timeout',
      '5000',
    ]);

    expect(mockRuntimeInstances[0]?.sendAgentMessage).toHaveBeenCalledWith({
      message: 'status update?',
      timeoutMs: 5000,
    });
    expect(errSpy).not.toHaveBeenCalled();
    logSpy.mockRestore();
    errSpy.mockRestore();
  });

  it('sets process exit code when options are invalid', async () => {
    const errSpy = vi.spyOn(console, 'error').mockImplementation(() => undefined);
    const program = new Command();
    const { registerCompanionCommand } = await import('./companion.js');
    registerCompanionCommand(program);

    await program.parseAsync(['node', 'test', 'companion', '--once', '--heartbeat', '0']);

    expect(errSpy).toHaveBeenCalled();
    expect(process.exitCode).toBe(1);

    errSpy.mockRestore();
  });
});