test: added tests for checking correct error codes are thrown on request timeout and connection closed scenarios

Author: sumnaith

src/shared/protocol.test.ts

@@ -0,0 +1,63 @@
+import { Protocol } from "./protocol.js";
+import { Transport } from "./transport.js";
+import {
+  McpError,
+  ErrorCode,
+  Request,
+  Result,
+  Notification,
+} from "../types.js";
+import { ZodType, z } from "zod";
+
+// Mock Transport class
+class MockTransport implements Transport {
+  onclose?: () => void;
+  onerror?: (error: Error) => void;
+  onmessage?: (message: unknown) => void;
+
+  async start(): Promise<void> {}
+  async close(): Promise<void> {
+    this.onclose?.();
+  }
+  async send(_message: unknown): Promise<void> {}
+}
+
+describe("protocol tests", () => {
+  let protocol: Protocol<Request, Notification, Result>;
+  let transport: MockTransport;
+
+  beforeEach(() => {
+    transport = new MockTransport();
+    protocol = new (class extends Protocol<Request, Notification, Result> {
+      protected assertCapabilityForMethod(): void {}
+      protected assertNotificationCapability(): void {}
+      protected assertRequestHandlerCapability(): void {}
+    })();
+  });
+
+  test("should throw a timeout error if the request exceeds the timeout", async () => {
+    await protocol.connect(transport);
+    const request = { method: "example", params: {} };
+    try {
+      const mockSchema: ZodType<{ result: string }> = z.object({
+        result: z.string(),
+      });
+      await protocol.request(request, mockSchema, {
+        timeout: 100,
+      }); // Short timeout for test
+    } catch (error) {
+      expect(error).toBeInstanceOf(McpError);
+      if (error instanceof McpError) {
+        expect(error.code).toBe(ErrorCode.RequestTimeout);
+      }
+    }
+  });
+
+  test("should invoke onclose when the connection is closed", async () => {
+    const oncloseMock = jest.fn();
+    protocol.onclose = oncloseMock;
+    await protocol.connect(transport);
+    await transport.close();
+    expect(oncloseMock).toHaveBeenCalled();
+  });
+});