feat: Added test for connect + improvements and fixes

Author: kevinwang5658

src/connect/login-helper.ts

@@ -65,6 +65,9 @@ export class LoginHelper {
       const credentialsPath = path.join(os.homedir(), '.codify', 'credentials.json');
       await fs.rm(credentialsPath);
     } catch {}
+
+    this.instance.isLoggedIn = false;
+    this.instance.credentials = undefined;
   }
 
   private static async read(): Promise<Credentials | undefined> {

src/orchestrators/connect.ts

@@ -14,12 +14,13 @@ import { SocketServer } from '../connect/socket-server.js';
 import { ProcessName, ctx } from '../events/context.js';
 import { Reporter } from '../ui/reporters/reporter.js';
 import { LoginOrchestrator } from './login.js';
+import { registerKillListeners } from '../utils/register-kill-listeners.js';
 
 export class ConnectOrchestrator {
   static rootCommand: string;
   static nodeBinary: string;
 
-  static async run(rootCommand: string, reporter: Reporter, openBrowser = true, onOpen?: (connectionCode: string) => void) {
+  static async run(rootCommand: string, reporter: Reporter, openBrowser = true, onOpen?: (connectionCode: string, server: Server) => void) {
     const login = LoginHelper.get()?.isLoggedIn;
     if (!login) {
       ctx.log('User is not logged in. Attempting to log in...')
@@ -36,7 +37,7 @@ export class ConnectOrchestrator {
     app.use(json())
     app.use(router);
 
-    const server = await ConnectOrchestrator.listen(app, reporter, () => {
+    const server = await ConnectOrchestrator.listen(app, reporter, (server) => {
       if (openBrowser) {
         open(`${config.dashboardUrl}/connection/success?code=${connectionSecret}`)
         console.log(`Open browser window to store code.
@@ -45,13 +46,13 @@ If unsuccessful manually enter the code:
 ${connectionSecret}`)
       }
 
-      onOpen?.(connectionSecret);
+      onOpen?.(connectionSecret, server);
     });
 
     SocketServer.init(server, connectionSecret);
   }
 
-  private static listen(app: express.Application, reporter: Reporter, onOpen: () => void): Promise<Server> {
+  private static listen(app: express.Application, reporter: Reporter, onOpen: (server: Server) => void): Promise<Server> {
      return new Promise((resolve) => {
       const server = app.listen(config.connectServerPort, async (error) => {
         if (error) {
@@ -77,9 +78,11 @@ ${connectionSecret}`)
           }
         } else {
           resolve(server);
-          onOpen();
+          onOpen(server);
         }
       });
+
+      registerKillListeners(() => server.close());
     });
   }
 

src/orchestrators/login.ts

@@ -7,6 +7,7 @@ import { DashboardApiClient } from '../api/dashboard/index.js';
 import { config } from '../config.js';
 import { LoginHelper } from '../connect/login-helper.js';
 import { ajv } from '../utils/ajv.js';
+import { registerKillListeners } from '../utils/register-kill-listeners.js';
 
 const schema = {
   type: 'object',
@@ -76,6 +77,8 @@ Manually open it here: ${config.dashboardUrl}/auth/cli`
       open(`${config.dashboardUrl}/auth/cli`);
     })
 
+    registerKillListeners(() => server.close());
+
     await new Promise<void>((resolve, reject) => {
       app.post('/', async (req, res) => {
         try {

src/plugins/plugin-manager.ts

@@ -11,6 +11,7 @@ import { ResourceConfig } from '../entities/resource-config.js';
 import { ResourceInfo } from '../entities/resource-info.js';
 import { SubProcessName, ctx } from '../events/context.js';
 import { groupBy } from '../utils/index.js';
+import { registerKillListeners } from '../utils/register-kill-listeners.js';
 import { Plugin } from './plugin.js';
 import { PluginResolver } from './resolver.js';
 
@@ -35,7 +36,11 @@ export class PluginManager {
       this.plugins.set(plugin.name, plugin)
     }
 
-    this.registerKillListeners(plugins)
+    registerKillListeners(() => {
+      for (const plugin of plugins) {
+        plugin.kill()
+      }
+    });
     return this.initializePlugins(plugins, secureMode, verbosityLevel);
   }
 
@@ -193,43 +198,4 @@ export class PluginManager {
 
     return resourceMap;
   }
-
-  /** Clean up any stranglers and child processes if the CLI is killed */
-  private registerKillListeners(plugins: Plugin[]) {
-    const kill = (code: number | string) => {
-      plugins.forEach((p) => {
-        p.kill()
-      })
-
-      let exitCode = 0;
-      switch (code) {
-        case 'SIGTERM': {
-          exitCode = 143;
-          break;
-        }
-
-        case 'SIGHUP': {
-          exitCode = 129;
-          break;
-        }
-
-        case 'SIGINT': {
-          exitCode = 130;
-          break;
-        }
-      }
-      
-      const parsedCode = typeof code === 'string' ? Number.parseInt(code, 10) : code;
-      if (Number.isInteger(parsedCode)) {
-        exitCode = parsedCode;
-      }
-
-      process.exit(exitCode);
-    }
-
-    process.on('exit', kill)
-    process.on('SIGINT', kill)
-    process.on('SIGTERM', kill)
-    process.on('SIGHUP', kill)
-  }
 }

src/utils/register-kill-listeners.ts

@@ -0,0 +1,35 @@
+export function registerKillListeners(kill: (code: number | string) => void) {
+  const killHandler = (code: number | string) => {
+    kill(code);
+
+    let exitCode = 0;
+    switch (code) {
+      case 'SIGTERM': {
+        exitCode = 143;
+        break;
+      }
+
+      case 'SIGHUP': {
+        exitCode = 129;
+        break;
+      }
+
+      case 'SIGINT': {
+        exitCode = 130;
+        break;
+      }
+    }
+
+    const parsedCode = typeof code === 'string' ? Number.parseInt(code, 10) : code;
+    if (Number.isInteger(parsedCode)) {
+      exitCode = parsedCode;
+    }
+
+    process.exit(exitCode);
+  }
+
+  process.on('exit', killHandler)
+  process.on('SIGINT', killHandler)
+  process.on('SIGTERM', killHandler)
+  process.on('SIGHUP', killHandler)
+}

test/orchestrator/connect/connect-command.test.ts

@@ -0,0 +1,116 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { MockOs } from '../mocks/system.js';
+import { MockReporter } from '../mocks/reporter.js';
+import { ConnectOrchestrator } from '../../../src/orchestrators/connect';
+import * as net from 'node:net';
+import { config } from '../../../src/config.js';
+import { fakeLogin, fakeLogout } from '../mocks/mock-login';
+import * as open from 'open'
+import { LoginOrchestrator } from '../../../src/orchestrators/login';
+import { Server } from 'node:http';
+import { vol } from 'memfs';
+
+vi.mock(import('../../../src/orchestrators/login'), async () => {
+  return {
+    LoginOrchestrator: {
+      run: async () => {}
+    },
+  }
+})
+
+vi.mock('node:fs', async () => {
+  const { fs } = await import('memfs');
+  return fs
+})
+
+vi.mock('node:fs/promises', async () => {
+  const { fs } = await import('memfs');
+  return fs.promises;
+})
+
+vi.mock(import('open'), async () => {
+  return {
+    default: vi.fn()
+  }
+})
+
+// The apply orchestrator directly calls plan so this will test both
+describe('Connect orchestrator tests', () => {
+  beforeEach(() => {
+    vol.reset();
+  })
+
+  it('It will start a local server on config.connectServerPort', async () => {
+    const reporter = new MockReporter();
+    await fakeLogin();
+
+    const openSpy = vi.spyOn(open, 'default');
+
+    await new Promise<void>((done) => {
+      ConnectOrchestrator.run('codify', reporter, false, async (connectionCode: string , server: Server) => {
+        expect(connectionCode).to.be.a('string');
+
+        const portInUse = await checkPortStatus(config.connectServerPort);
+        expect(portInUse).to.be.true;
+
+        server.close();
+        done();
+      })
+    })
+  });
+
+  it('It will ask for a login if the user is not logged in', async () => {
+    const reporter = new MockReporter({});
+    await fakeLogout();
+
+    const loginRunSpy = vi.spyOn(LoginOrchestrator, 'run');
+    const openSpy = vi.spyOn(open, 'default');
+
+    await new Promise<void>((done) => {
+      ConnectOrchestrator.run('codify', reporter, false, async (connectionCode: string, server: Server) => {
+        expect(connectionCode).to.be.a('string');
+
+        const portInUse = await checkPortStatus(config.connectServerPort);
+        expect(portInUse).to.be.true;
+        expect(loginRunSpy).toHaveBeenCalledOnce();
+
+        done();
+        server.close();
+      })
+
+    });
+  });
+
+  const wait = (ms: number) => new Promise(resolve => setTimeout(resolve, ms));
+
+  afterEach(() => {
+    vi.resetAllMocks();
+    MockOs.reset();
+  })
+
+  function checkPortStatus(port: number, host = '127.0.0.1') {
+    return new Promise((resolve, reject) => {
+      const socket = new net.Socket();
+
+      socket.once('connect', () => {
+        // If 'connect' event fires, the port is open and listening
+        socket.destroy();
+        resolve(true); // Port is in use
+      });
+
+      socket.once('error', (err) => {
+        // Any error typically means the port is not listening
+        // EADDRNOTAVAIL, ECONNREFUSED, etc.
+        reject(err); // Port is likely free or unreachable
+      });
+
+      socket.once('timeout', () => {
+        socket.destroy();
+        reject(new Error('Connection attempt timed out'));
+      });
+
+      socket.connect(port, host);
+    });
+  }
+
+})

test/orchestrator/mocks/mock-login.ts

@@ -0,0 +1,21 @@
+import { vi } from 'vitest';
+
+import { LoginHelper } from '../../../src/connect/login-helper.js';
+
+/**
+ * Must mock node:fs/promises before calling this function
+ */
+export async function fakeLogin(): Promise<string> {
+  await LoginHelper.load();
+  const token = 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWUsImlhdCI6MTUxNjIzOTAyMn0.KMUFsIDTnFmyG3nMiGM6H9FNFUROf3wh7SmqJp-QV30'
+  await LoginHelper.save(token);
+  return token;
+}
+
+/**
+ * Must mock node:fs/promises before calling this function
+ */
+export async function fakeLogout() {
+  await LoginHelper.load();
+  await LoginHelper.logout();
+}