Yonexus/tests/integration/framework.test.ts

import { describe, expect, it, vi, beforeEach, afterEach } from "vitest";
import type { ClientConnection, ServerTransport } from "../Yonexus.Server/plugin/core/transport.js";
import type { ClientTransport } from "../Yonexus.Client/plugin/core/transport.js";
import type { YonexusServerStore } from "../Yonexus.Server/plugin/core/store.js";
import type { YonexusClientStateStore } from "../Yonexus.Client/plugin/core/state.js";
import { createYonexusServerRuntime } from "../Yonexus.Server/plugin/core/runtime.js";
import { createYonexusClientRuntime } from "../Yonexus.Client/plugin/core/runtime.js";
import {
  decodeBuiltin,
  encodeBuiltin,
  buildHello,
  buildHelloAck,
  buildPairRequest,
  buildPairConfirm,
  buildPairSuccess,
  buildAuthRequest,
  buildAuthSuccess,
  buildHeartbeat,
  buildHeartbeatAck,
  YONEXUS_PROTOCOL_VERSION
} from "../Yonexus.Protocol/src/index.js";
import { generateKeyPair } from "../Yonexus.Client/plugin/crypto/keypair.js";
import type { ClientRecord } from "../Yonexus.Server/plugin/core/persistence.js";
import type { YonexusClientState } from "../Yonexus.Client/plugin/core/state.js";

/**
 * Yonexus Server-Client Integration Test Framework
 *
 * This module provides utilities for testing Server and Client interactions
 * without requiring real network sockets.
 */

// ============================================================================
// Mock Transport Pair - Simulates network connection between Server and Client
// ============================================================================

export interface MockMessageChannel {
  serverToClient: string[];
  clientToServer: string[];
}

export interface MockTransportPair {
  serverTransport: ServerTransport;
  clientTransport: ClientTransport;
  channel: MockMessageChannel;
  getServerReceived: () => string[];
  getClientReceived: () => string[];
  clearMessages: () => void;
}

export function createMockTransportPair(): MockTransportPair {
  const channel: MockMessageChannel = {
    serverToClient: [],
    clientToServer: []
  };

  // Track server-side connections
  const serverConnections = new Map<string, ClientConnection>();
  let tempConnection: ClientConnection | null = null;

  // Server Transport Mock
  const serverTransport: ServerTransport = {
    isRunning: false,
    connections: serverConnections,

    start: vi.fn(async () => {
      serverTransport.isRunning = true;
    }),

    stop: vi.fn(async () => {
      serverTransport.isRunning = false;
      serverConnections.clear();
    }),

    send: vi.fn((identifier: string, message: string) => {
      if (serverConnections.has(identifier)) {
        channel.serverToClient.push(message);
        return true;
      }
      return false;
    }),

    sendToConnection: vi.fn((connection: ClientConnection, message: string) => {
      channel.serverToClient.push(message);
      return true;
    }),

    broadcast: vi.fn((message: string) => {
      channel.serverToClient.push(`[broadcast]:${message}`);
    }),

    closeConnection: vi.fn((identifier: string, code?: number, reason?: string) => {
      const conn = serverConnections.get(identifier);
      if (conn) {
        conn.isAuthenticated = false;
        serverConnections.delete(identifier);
      }
      return true;
    }),

    assignIdentifierToTemp: vi.fn((ws, identifier: string) => {
      if (tempConnection) {
        tempConnection.identifier = identifier;
      }
    }),

    promoteToAuthenticated: vi.fn((identifier: string, ws) => {
      if (tempConnection && tempConnection.identifier === identifier) {
        tempConnection.isAuthenticated = true;
        serverConnections.set(identifier, tempConnection);
        tempConnection = null;
      }
    }),

    removeTempConnection: vi.fn(() => {
      tempConnection = null;
    })
  };

  // Client Transport Mock
  let clientState: import("../Yonexus.Client/plugin/core/transport.js").ClientConnectionState = "idle";

  const clientTransport: ClientTransport = {
    get state() {
      return clientState;
    },

    get isConnected() {
      return clientState !== "idle" && clientState !== "disconnected" && clientState !== "error";
    },

    get isAuthenticated() {
      return clientState === "authenticated";
    },

    connect: vi.fn(async () => {
      clientState = "connected";
      // Simulate connection - create temp connection on server side
      tempConnection = {
        identifier: null,
        ws: { close: vi.fn() } as unknown as WebSocket,
        connectedAt: Date.now(),
        isAuthenticated: false
      };
    }),

    disconnect: vi.fn(() => {
      clientState = "disconnected";
      tempConnection = null;
    }),

    send: vi.fn((message: string) => {
      if (clientState === "connected" || clientState === "authenticated" || clientState === "authenticating") {
        channel.clientToServer.push(message);
        return true;
      }
      return false;
    }),

    markAuthenticated: vi.fn(() => {
      clientState = "authenticated";
    }),

    markAuthenticating: vi.fn(() => {
      clientState = "authenticating";
    })
  };

  return {
    serverTransport,
    clientTransport,
    channel,
    getServerReceived: () => [...channel.clientToServer],
    getClientReceived: () => [...channel.serverToClient],
    clearMessages: () => {
      channel.serverToClient.length = 0;
      channel.clientToServer.length = 0;
    }
  };
}

// ============================================================================
// Mock Store Factories
// ============================================================================

export function createMockServerStore(initialClients: ClientRecord[] = []): YonexusServerStore {
  const persisted = new Map(initialClients.map((record) => [record.identifier, record]));

  return {
    filePath: "/tmp/yonexus-server-test.json",
    load: vi.fn(async () => ({
      version: 1,
      persistedAt: Date.now(),
      clients: new Map(persisted)
    })),
    save: vi.fn(async (clients: Iterable<ClientRecord>) => {
      persisted.clear();
      for (const client of clients) {
        persisted.set(client.identifier, client);
      }
    })
  };
}

export function createMockClientStore(initialState?: Partial<YonexusClientState>): YonexusClientStateStore {
  let state: YonexusClientState = {
    identifier: initialState?.identifier ?? "test-client",
    publicKey: initialState?.publicKey,
    privateKey: initialState?.privateKey,
    secret: initialState?.secret,
    pairedAt: initialState?.pairedAt,
    authenticatedAt: initialState?.authenticatedAt,
    updatedAt: initialState?.updatedAt ?? Date.now()
  };

  return {
    filePath: "/tmp/yonexus-client-test.json",
    load: vi.fn(async () => ({ ...state })),
    save: vi.fn(async (next) => {
      state = { ...next };
    })
  };
}

// ============================================================================
// Test Runtime Factory
// ============================================================================

export interface IntegrationTestContext {
  serverRuntime: ReturnType<typeof createYonexusServerRuntime>;
  clientRuntime: ReturnType<typeof createYonexusClientRuntime>;
  transports: MockTransportPair;
  serverStore: YonexusServerStore;
  clientStore: YonexusClientStateStore;
  advanceTime: (seconds: number) => void;
  processServerToClient: () => Promise<void>;
  processClientToServer: () => Promise<void>;
  processAllMessages: () => Promise<void>;
}

export async function createIntegrationTestContext(
  options: {
    clientIdentifier?: string;
    paired?: boolean;
    authenticated?: boolean;
    serverTime?: number;
  } = {}
): Promise<IntegrationTestContext> {
  const now = options.serverTime ?? 1_710_000_000;
  const identifier = options.clientIdentifier ?? "test-client";

  const transports = createMockTransportPair();
  const serverStore = createMockServerStore();
  const clientStore = createMockClientStore({ identifier });

  // Generate keypair for client if needed
  const keyPair = await generateKeyPair();

  const serverRuntime = createYonexusServerRuntime({
    config: {
      followerIdentifiers: [identifier],
      notifyBotToken: "test-token",
      adminUserId: "admin-user",
      listenHost: "127.0.0.1",
      listenPort: 8787
    },
    store: serverStore,
    transport: transports.serverTransport,
    now: () => now
  });

  const clientRuntime = createYonexusClientRuntime({
    config: {
      mainHost: "ws://localhost:8787",
      identifier,
      notifyBotToken: "test-token",
      adminUserId: "admin-user"
    },
    transport: transports.clientTransport,
    stateStore: clientStore,
    now: () => now
  });

  await serverRuntime.start();

  let currentTime = now;
  const advanceTime = (seconds: number) => {
    currentTime += seconds;
  };

  // Message processing helpers
  const processServerToClient = async () => {
    const messages = transports.getClientReceived();
    transports.clearMessages();
    for (const msg of messages) {
      await clientRuntime.handleMessage(msg);
    }
  };

  const processClientToServer = async () => {
    const messages = transports.getServerReceived();
    transports.clearMessages();

    // Get the temp connection for message handling
    const connection = {
      identifier: identifier,
      ws: { close: vi.fn() } as unknown as WebSocket,
      connectedAt: currentTime,
      isAuthenticated: options.authenticated ?? false
    };

    for (const msg of messages) {
      await serverRuntime.handleMessage(connection, msg);
    }
  };

  const processAllMessages = async () => {
    await processClientToServer();
    await processServerToClient();
  };

  return {
    serverRuntime,
    clientRuntime,
    transports,
    serverStore,
    clientStore,
    advanceTime,
    processServerToClient,
    processClientToServer,
    processAllMessages
  };
}

// ============================================================================
// Integration Test Suite
// ============================================================================

describe("Yonexus Server-Client Integration", () => {
  beforeEach(() => {
    vi.useFakeTimers();
  });

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

  describe("First-Time Pairing Flow", () => {
    it("completes full pairing and authentication cycle", async () => {
      const ctx = await createIntegrationTestContext({
        clientIdentifier: "new-client"
      });

      // Step 1: Client connects and sends hello
      await ctx.clientRuntime.start();
      ctx.clientRuntime.handleTransportStateChange("connected");
      await vi.advanceTimersByTimeAsync(100);

      // Process hello -> hello_ack + pair_request
      await ctx.processClientToServer();
      await ctx.processServerToClient();

      // Verify client received pair_request
      expect(ctx.clientRuntime.state.phase).toBe("waiting_pair_confirm");
      expect(ctx.clientRuntime.state.pendingPairing).toBeDefined();

      // Step 2: Client submits pairing code
      const pairingCode = ctx.serverRuntime.state.registry.clients.get("new-client")?.pairingCode;
      expect(pairingCode).toBeDefined();

      ctx.clientRuntime.submitPairingCode(pairingCode!, "req-pair-confirm");
      await vi.advanceTimersByTimeAsync(100);

      // Process pair_confirm -> pair_success
      await ctx.processClientToServer();
      await ctx.processServerToClient();

      // Verify client received secret
      expect(ctx.clientRuntime.state.clientState.secret).toBeDefined();
      expect(ctx.clientRuntime.state.phase).toBe("auth_required");

      // Step 3: Client sends auth request
      await vi.advanceTimersByTimeAsync(100);
      await ctx.processClientToServer();
      await ctx.processServerToClient();

      // Verify authentication success
      expect(ctx.clientRuntime.state.phase).toBe("authenticated");
      expect(ctx.serverRuntime.state.registry.sessions.get("new-client")?.isAuthenticated).toBe(true);
    });
  });

  describe("Reconnection Flow", () => {
    it("reconnects with existing credentials without re-pairing", async () => {
      const now = 1_710_000_000;
      const keyPair = await generateKeyPair();

      const ctx = await createIntegrationTestContext({
        clientIdentifier: "reconnect-client",
        paired: true,
        authenticated: false
      });

      // Pre-populate client with existing credentials
      ctx.clientRuntime.state.clientState.secret = "existing-secret";
      ctx.clientRuntime.state.clientState.publicKey = keyPair.publicKey;
      ctx.clientRuntime.state.clientState.privateKey = keyPair.privateKey;
      ctx.clientRuntime.state.clientState.pairedAt = now - 1000;

      // Pre-populate server with client record
      ctx.serverRuntime.state.registry.clients.set("reconnect-client", {
        identifier: "reconnect-client",
        pairingStatus: "paired",
        publicKey: keyPair.publicKey,
        secret: "existing-secret",
        status: "offline",
        recentNonces: [],
        recentHandshakeAttempts: [],
        createdAt: now - 2000,
        updatedAt: now - 1000
      });

      // Connect and send hello
      await ctx.clientRuntime.start();
      ctx.clientRuntime.handleTransportStateChange("connected");
      await vi.advanceTimersByTimeAsync(100);

      await ctx.processClientToServer();
      await ctx.processServerToClient();

      // Should go directly to auth_required, skipping pairing
      expect(ctx.clientRuntime.state.phase).toBe("auth_required");

      // Complete authentication
      await vi.advanceTimersByTimeAsync(100);
      await ctx.processClientToServer();
      await ctx.processServerToClient();

      expect(ctx.clientRuntime.state.phase).toBe("authenticated");
    });
  });

  describe("Heartbeat Flow", () => {
    it("exchanges heartbeats after authentication", async () => {
      const ctx = await createIntegrationTestContext({
        clientIdentifier: "heartbeat-client",
        paired: true,
        authenticated: true
      });

      // Manually set authenticated state
      ctx.clientRuntime.state.phase = "authenticated";

      // Trigger heartbeat
      await ctx.clientRuntime.handleMessage("heartbeat_tick");
      await vi.advanceTimersByTimeAsync(100);

      // Process heartbeat -> heartbeat_ack
      await ctx.processClientToServer();

      // Verify server updated heartbeat timestamp
      const record = ctx.serverRuntime.state.registry.clients.get("heartbeat-client");
      expect(record?.lastHeartbeatAt).toBeDefined();
    });
  });
});