From 9bd62e5ee98e9248340a4423ab0a6b319b5bd6a7 Mon Sep 17 00:00:00 2001
From: nav <nav@mgr.hangman-lab.top>
Date: Thu, 9 Apr 2026 02:04:06 +0000
Subject: [PATCH] test: cover connection failure edge cases

---
 tests/connection-heartbeat-failures.test.ts | 111 +++++++++++++++++++-
 1 file changed, 106 insertions(+), 5 deletions(-)

diff --git a/tests/connection-heartbeat-failures.test.ts b/tests/connection-heartbeat-failures.test.ts
index 659e16f..e9e43f3 100644
--- a/tests/connection-heartbeat-failures.test.ts
+++ b/tests/connection-heartbeat-failures.test.ts
@@ -1,6 +1,12 @@
 import { describe, expect, it, vi, beforeEach, afterEach } from "vitest";
 
-import { buildHeartbeat, decodeBuiltin, encodeBuiltin } from "../../Yonexus.Protocol/src/index.js";
+import {
+  buildHeartbeat,
+  buildHello,
+  decodeBuiltin,
+  encodeBuiltin,
+  YONEXUS_PROTOCOL_VERSION
+} from "../../Yonexus.Protocol/src/index.js";
 import { createYonexusServerRuntime } from "../plugin/core/runtime.js";
 import { createClientRecord, type ClientRecord } from "../plugin/core/persistence.js";
 import type { YonexusServerStore } from "../plugin/core/store.js";
@@ -40,10 +46,12 @@ function createMockStore(initialClients: ClientRecord[] = []): YonexusServerStor
 
 function createMockTransport() {
   const sent: Array<{ connection: ClientConnection; message: string }> = [];
+  const tempAssignments = new Map<ClientConnection["ws"], string>();
+  const connections = new Map<string, ClientConnection>();
 
   const transport: ServerTransport = {
     isRunning: false,
-    connections: new Map(),
+    connections,
     start: vi.fn(),
     stop: vi.fn(),
     send: vi.fn(),
@@ -53,9 +61,31 @@ function createMockTransport() {
     }),
     broadcast: vi.fn(),
     closeConnection: vi.fn(),
-    promoteToAuthenticated: vi.fn(),
-    removeTempConnection: vi.fn(),
-    assignIdentifierToTemp: vi.fn()
+    promoteToAuthenticated: vi.fn((identifier: string, ws: ClientConnection["ws"]) => {
+      if (!tempAssignments.has(ws)) {
+        return false;
+      }
+
+      const existing = connections.get(identifier);
+      if (existing) {
+        existing.ws.close(1008, "Connection replaced by new authenticated session");
+      }
+
+      connections.set(identifier, {
+        identifier,
+        ws,
+        connectedAt: 1_710_000_000,
+        isAuthenticated: true
+      });
+      tempAssignments.delete(ws);
+      return true;
+    }),
+    removeTempConnection: vi.fn((ws: ClientConnection["ws"]) => {
+      tempAssignments.delete(ws);
+    }),
+    assignIdentifierToTemp: vi.fn((ws: ClientConnection["ws"], identifier: string) => {
+      tempAssignments.set(ws, identifier);
+    })
   };
 
   return { transport, sent };
@@ -193,4 +223,75 @@ describe("YNX-1105d: Connection & Heartbeat Failure Paths", () => {
       code: "AUTH_FAILED"
     });
   });
+
+  it("CF-04: protocol version mismatch returns error and closes the connection", async () => {
+    const record = createClientRecord("client-a");
+    const store = createMockStore([record]);
+    const { transport, sent } = createMockTransport();
+
+    const runtime = createYonexusServerRuntime({
+      config: {
+        followerIdentifiers: ["client-a"],
+        notifyBotToken: "stub-token",
+        adminUserId: "admin-user",
+        listenHost: "127.0.0.1",
+        listenPort: 8787
+      },
+      store,
+      transport,
+      now: () => now
+    });
+
+    await runtime.start();
+
+    const connection = createConnection();
+    await runtime.handleMessage(
+      connection,
+      encodeBuiltin(
+        buildHello(
+          {
+            identifier: "client-a",
+            hasSecret: false,
+            hasKeyPair: false,
+            protocolVersion: `${YONEXUS_PROTOCOL_VERSION}-unsupported`
+          },
+          { requestId: "req-hello-version", timestamp: now }
+        )
+      )
+    );
+
+    const lastMessage = decodeBuiltin(sent.at(-1)!.message);
+    expect(lastMessage.type).toBe("error");
+    expect(lastMessage.payload).toMatchObject({
+      code: "UNSUPPORTED_PROTOCOL_VERSION"
+    });
+    expect(connection.ws.close).toHaveBeenCalledWith(1002, "Unsupported protocol version");
+  });
+
+  it("CF-03: promoting a new authenticated connection replaces the old one", async () => {
+    const previousSocket = createMockSocket();
+    const replacementSocket = createMockSocket();
+    const previousConnection = {
+      identifier: "client-a",
+      ws: previousSocket,
+      connectedAt: now - 5,
+      isAuthenticated: true
+    } satisfies ClientConnection;
+
+    const { transport } = createMockTransport();
+    transport.connections.set("client-a", previousConnection);
+
+    transport.assignIdentifierToTemp(replacementSocket, "client-a");
+    const promoted = transport.promoteToAuthenticated("client-a", replacementSocket);
+
+    expect(promoted).toBe(true);
+    expect(previousSocket.close).toHaveBeenCalledWith(
+      1008,
+      "Connection replaced by new authenticated session"
+    );
+
+    const activeConnection = transport.connections.get("client-a");
+    expect(activeConnection?.ws).toBe(replacementSocket);
+    expect(activeConnection?.isAuthenticated).toBe(true);
+  });
 });