Fix the flaky worker integration tests

integration-tests/worker/src/init.ts

@@ -3,14 +3,16 @@ import crypto from 'node:crypto';
 
 import createLightningServer, { toBase64 } from '@openfn/lightning-mock';
 import createEngine from '@openfn/engine-multi';
-import createWorkerServer from '@openfn/ws-worker';
+import createWorkerServer, { INTERNAL_SOCKET_READY } from '@openfn/ws-worker';
 import { createMockLogger } from '@openfn/logger';
 import createLogger from '@openfn/logger';
 
 const debugWorker = process.env.OPENFN_DEBUG_WORKER;
 const debugLightning = process.env.OPENFN_DEBUG_LIGHTNING;
 
-export const randomPort = () => Math.round(2000 + Math.random() * 1000);
+let nextPortNumber = 4400;
+
+export const nextPort = () => nextPortNumber++;
 
 export const initLightning = (port = 4000, privateKey?: string) => {
   // TODO the lightning mock right now doesn't use the secret
@@ -30,7 +32,7 @@ export const initWorker = async (
   engineArgs = {},
   workerArgs = {}
 ) => {
-  const workerPort = randomPort();
+  const workerPort = nextPort();
 
   const engineLogger = createMockLogger('engine', {
     level: 'debug',
@@ -56,5 +58,18 @@ export const initWorker = async (
     ...workerArgs,
   });
 
+  if (!worker.socket) {
+    await new Promise<void>((resolve, reject) => {
+      const timeout = setTimeout(() => {
+        reject(new Error('worker did not connect to lightning within 10s'));
+      }, 10_000);
+
+      worker.events.once(INTERNAL_SOCKET_READY, () => {
+        clearTimeout(timeout);
+        resolve();
+      });
+    });
+  }
+
   return { engine, engineLogger, worker };
 };

integration-tests/worker/test/autoinstall.test.ts

@@ -3,7 +3,7 @@ import path from 'node:path';
 import { rm } from 'node:fs/promises';
 import { generateKeys } from '@openfn/lightning-mock';
 
-import { initLightning, initWorker } from '../src/init';
+import { initLightning, initWorker, nextPort } from '../src/init';
 import { createRun, createJob } from '../src/factories';
 
 const generate = (adaptor, version) => {
@@ -39,7 +39,7 @@ test.before(async () => {
   } catch (e) {}
 
   const keys = await generateKeys();
-  const lightningPort = 4321;
+  const lightningPort = nextPort();
 
   lightning = initLightning(lightningPort, keys.private);
 

integration-tests/worker/test/benchmark.test.ts

@@ -2,7 +2,7 @@ import test from 'ava';
 import path from 'node:path';
 
 import { createRun } from '../src/factories';
-import { initLightning, initWorker } from '../src/init';
+import { initLightning, initWorker, nextPort } from '../src/init';
 import { run, humanMb } from '../src/util';
 
 let lightning;
@@ -11,7 +11,7 @@ let worker;
 const maxConcurrency = 20;
 
 test.before(async () => {
-  const lightningPort = 4322;
+  const lightningPort = nextPort();
 
   lightning = initLightning(lightningPort);
 

integration-tests/worker/test/exit-reasons.test.ts

@@ -2,13 +2,13 @@ import test from 'ava';
 import crypto from 'node:crypto';
 import path from 'node:path';
 
-import { initLightning, initWorker } from '../src/init';
+import { initLightning, initWorker, nextPort } from '../src/init';
 
 let lightning;
 let worker;
 
 test.before(async () => {
-  const lightningPort = 4321;
+  const lightningPort = nextPort();
 
   lightning = initLightning(lightningPort);
 
@@ -23,8 +23,8 @@ test.after(async () => {
 });
 
 const run = async (attempt) => {
-  return new Promise<any>(async (done) => {
-    lightning.once('run:complete', (evt) => {
+  return new Promise<any>((done) => {
+    lightning.onSocketEvent('run:complete', attempt.id, (evt) => {
       if (attempt.id === evt.runId) {
         done(evt.payload);
       }

integration-tests/worker/test/integration.test.ts

@@ -4,7 +4,7 @@ import crypto from 'node:crypto';
 import Koa from 'koa';
 import { generateKeys } from '@openfn/lightning-mock';
 
-import { initLightning, initWorker, randomPort } from '../src/init';
+import { initLightning, initWorker, nextPort } from '../src/init';
 
 let lightning;
 let worker;
@@ -14,7 +14,7 @@ let lightningPort;
 
 test.before(async () => {
   const keys = await generateKeys();
-  lightningPort = randomPort();
+  lightningPort = nextPort();
   lightning = initLightning(lightningPort, keys.private);
 
   const engineArgs = {

integration-tests/worker/test/runs.test.ts

@@ -8,14 +8,14 @@ import {
   createJob,
   createTrigger,
 } from '../src/factories';
-import { initLightning, initWorker } from '../src/init';
+import { initLightning, initWorker, nextPort } from '../src/init';
 
 let lightning;
 let worker;
 
 test.before(async () => {
   const keys = await generateKeys();
-  const lightningPort = 4321;
+  const lightningPort = nextPort();
 
   lightning = initLightning(lightningPort, keys.private);
 

integration-tests/worker/test/server.test.ts

@@ -180,7 +180,6 @@ test.serial('allow a job to complete after receiving a sigterm', (t) => {
 test.serial("don't restore the claim loop after a sigterm", (t) => {
   return new Promise(async (done) => {
     let abort = false;
-    let didSendSigterm = false;
     const port = getPort();
 
     const job = createJob({
@@ -208,7 +207,6 @@ test.serial("don't restore the claim loop after a sigterm", (t) => {
     // After the second run starts, there should be no more claims
     lightning.on('run:start', (evt) => {
       if (evt.runId === b.id) {
-        didSendSigterm = true;
         // Kill the worker once the second job has started
         // This will force an overlap of two pending runs at full capacity
         workerProcess.kill('SIGTERM');
@@ -220,8 +218,12 @@ test.serial("don't restore the claim loop after a sigterm", (t) => {
     lightning.enqueueRun(a);
     lightning.enqueueRun(b);
 
-    lightning.on('claim', (e) => {
-      if (didSendSigterm && !abort) {
+    lightning.on('claim', () => {
+      const didReceiveSigterm = workerLogs.some((l) =>
+        l.match(/SIGTERM RECEIVED/)
+      );
+
+      if (didReceiveSigterm && !abort) {
         abort = true;
         t.fail('Claim triggered after sigterm');
         done();