test: Improve server mock and assertions

Author: alvinveroy

src/lib/server.ts

@@ -1690,7 +1690,7 @@ export async function startProxyServer(
     // const tempServerForPortCheck = http.createServer(); // Not needed if findFreePort creates its own
     proxyListenPort = await findFreePort(initialPortForProxySearch); // findFreePort internally creates server for checks
     logger.info(`[PROXY_DEBUG] startProxyServer: Found free port ${proxyListenPort} for proxy.`);
-  } catch (error) { // Keep error as any or unknown for broader catch
+  } catch (error: unknown) { // Explicitly type error as unknown
     const err = error instanceof Error ? error : new Error(String(error));
     logger.error(`[ProxyServer] Failed to find free port for proxy: ${err.message}`, { errorDetails: err });
     return null; // Return null if findFreePort fails

src/tests/index.test.ts

@@ -568,6 +568,7 @@ describe('CLI with yargs (index.ts)', () => {
       mockConsoleError.mockClear(); // Clear before run
 
       await expect(runMainWithArgs(['agent_query', '{"query": "test"'])).rejects.toThrowError(
+        // This error is thrown by handleClientCommand and then re-thrown by yargs.fail()
         new Error("Invalid JSON parameters: Expected ',' or '}' after property value in JSON at position 16")
       );
 
@@ -659,11 +660,10 @@ describe('CLI with yargs (index.ts)', () => {
       expect(mockStartServerHandler).toHaveBeenCalled(); 
 
       // Check if the SUT's yargs middleware logged the port being observed.
-      expect(currentMockLoggerInstance.info.mock.calls).toEqual(
-        expect.arrayContaining([
-          [expect.stringContaining(`[SUT_INDEX_TS_YARGS_MW] --port option value at middleware: ${customPort}`)]
-        ])
+      const portLogFound = currentMockLoggerInstance.info.mock.calls.some(
+        call => typeof call[0] === 'string' && call[0].includes(`[SUT_INDEX_TS_YARGS_MW] --port option value at middleware: ${customPort}`)
       );
+      expect(portLogFound, `Expected SUT log for --port middleware. Calls: ${JSON.stringify(currentMockLoggerInstance.info.mock.calls)}`).toBe(true);
 
       // Restore original
       if (originalHttpPort === undefined) delete process.env.HTTP_PORT;
@@ -778,7 +778,7 @@ describe('CLI with yargs (index.ts)', () => {
       mockConsoleError.mockClear(); // Clear before run
       mockProcessExit.mockClear();
 
-      const expectedErrorPattern = /Unknown command: unknowncommand|You must provide a command to run|Not enough non-option arguments|Unknown argument: unknowncommand/; // Added Unknown argument
+      const expectedErrorPattern = /Unknown command: unknowncommand|You must provide a command to run|Not enough non-option arguments|Unknown argument: unknowncommand/;
       await expect(runMainWithArgs(['unknowncommand'])).rejects.toThrowError(expectedErrorPattern);
 
       const consoleErrorCalls = mockConsoleError.mock.calls.map(call => call.join(' ')).join('\n');
@@ -881,10 +881,15 @@ describe('CLI with yargs (index.ts)', () => {
 
       process.env.VITEST_TESTING_FAIL_HANDLER = "true";
       await expect(runMainWithArgs(['agent_query', '{"query":"test_json_rpc_error"}', '--json']))
-        .rejects.toThrow(expect.objectContaining({ // Check for an object that contains the jsonRpcError
-          message: `Tool 'agent_query' failed: ${rpcErrorObject.message}`,
-          jsonRpcError: rpcErrorObject
-        }));
+        .rejects.toThrow(
+          // The error thrown by yargs.fail() will be an Error instance.
+          // If handleClientCommand throws an error with a jsonRpcError property,
+          // yargs.fail() should re-throw that error object.
+          expect.objectContaining({
+            message: `Tool 'agent_query' failed: ${rpcErrorObject.message}`,
+            jsonRpcError: rpcErrorObject,
+          })
+        );
 
       // SUT's handleClientCommand logs the JSON error to console.error when --json is active
       // It logs the error object passed to it, which now includes jsonRpcError
@@ -906,7 +911,7 @@ describe('CLI with yargs (index.ts)', () => {
         call[0] === 'YARGS_FAIL_TEST_MODE_ERROR_OUTPUT:' && 
         typeof call[1] === 'string' && call[1].includes(`Tool 'agent_query' failed: ${rpcErrorObject.message}`)
       );
-      expect(yargsFailOutputLogged).toBe(true);
+      expect(yargsFailOutputLogged, `Expected yargs fail output. Calls: ${JSON.stringify(mockConsoleError.mock.calls)}`).toBe(true);
       expect(mockProcessExit).not.toHaveBeenCalled();
       delete process.env.VITEST_TESTING_FAIL_HANDLER;
     });
@@ -935,7 +940,7 @@ describe('CLI with yargs (index.ts)', () => {
         call[0] === 'YARGS_FAIL_TEST_MODE_ERROR_OUTPUT:' && 
         typeof call[1] === 'string' && call[1] === genericError.message
       );
-      expect(yargsFailOutputLogged).toBe(true);
+      expect(yargsFailOutputLogged, `Expected yargs fail output. Calls: ${JSON.stringify(mockConsoleError.mock.calls)}`).toBe(true);
       expect(mockProcessExit).not.toHaveBeenCalled();
       delete process.env.VITEST_TESTING_FAIL_HANDLER;
     });

src/tests/server.test.ts

@@ -1385,11 +1385,11 @@ describe('startProxyServer', () => {
           this._host = hostToListen ?? null;
 
           process.nextTick(() => {
-            if (this._listeners && typeof this._listeners.listening === 'function') {
-              (this._listeners.listening as () => void)();
-            } else if (Array.isArray(this._listeners?.listening)) {
-                (this._listeners.listening as ((...args: any[]) => void)[]).forEach(fn => fn());
+            // Emit 'listening' event
+            if (typeof this.emit === 'function') {
+              this.emit('listening');
             }
+            // Call direct callback if provided
             if (actualCallback) actualCallback();
           });
           return this; 
@@ -1398,23 +1398,23 @@ describe('startProxyServer', () => {
           if (!this._listeners[event]) {
             this._listeners[event] = [];
           }
-          if (Array.isArray(this._listeners[event])) {
-            (this._listeners[event] as ((...args: any[]) => void)[]).push(callback);
-          } else { // If it was a single function (e.g. from 'once'), make it an array
-            this._listeners[event] = [this._listeners[event] as (...args: any[]) => void, callback];
+          const listenersArray = this._listeners[event] as ((...args: any[]) => void)[];
+          if (!Array.isArray(listenersArray)) { // Should not happen if initialized as array
+            this._listeners[event] = [callback];
+          } else {
+            listenersArray.push(callback);
           }
           return this;
         }),
         once: vi.fn(function(this: any, event: string, callback: (...args: any[]) => void) {
           const onceWrapper = (...args: any[]) => {
-            if (Array.isArray(this._listeners[event])) {
-              this._listeners[event] = (this._listeners[event] as ((...args: any[]) => void)[]).filter(fn => fn !== onceWrapper);
-            } else if (this._listeners[event] === onceWrapper) {
-              delete this._listeners[event];
+            const listenersArray = this._listeners[event] as ((...args: any[]) => void)[];
+            if (Array.isArray(listenersArray)) {
+              this._listeners[event] = listenersArray.filter(fn => fn !== onceWrapper);
             }
             callback(...args);
           };
-          this.on(event, onceWrapper); // Use the 'on' method to add it
+          this.on(event, onceWrapper);
           return this;
         }),
         address: vi.fn(function(this: any) {
@@ -1425,10 +1425,8 @@ describe('startProxyServer', () => {
         }),
         close: vi.fn(function(this: any, cb?: (err?: Error) => void) { 
           process.nextTick(() => {
-            if (this._listeners && typeof this._listeners.close === 'function') {
-                (this._listeners.close as () => void)();
-            } else if (Array.isArray(this._listeners?.close)) {
-                (this._listeners.close as ((...args: any[]) => void)[]).forEach(fn => fn());
+            if (typeof this.emit === 'function') {
+              this.emit('close');
             }
             if (cb) cb();
           });
@@ -1437,10 +1435,11 @@ describe('startProxyServer', () => {
         removeAllListeners: vi.fn().mockReturnThis(),
         emit: vi.fn(function(this: any, event: string, ...args: any[]) {
             const listeners = this._listeners?.[event];
-            if (typeof listeners === 'function') {
+            if (typeof listeners === 'function') { // Should not happen with 'on' pushing to array
                 listeners(...args);
             } else if (Array.isArray(listeners)) {
-                listeners.forEach(fn => fn(...args));
+                // Iterate over a copy in case a listener removes itself
+                [...listeners].forEach(fn => fn(...args));
             }
         }),
       };