Fix more tests

Author: ericallam

internal-packages/run-engine/src/engine/systems/dequeueSystem.ts

@@ -542,84 +542,84 @@ export class DequeueSystem {
   }) {
     const prisma = tx ?? this.$.prisma;
 
-    return startSpan(
-      this.$.tracer,
-      "#pendingVersion",
-      async (span) => {
-        return this.$.runLock.lock("pendingVersion", [runId], 5_000, async (signal) => {
-          //mark run as waiting for deploy
-          const run = await prisma.taskRun.update({
-            where: { id: runId },
-            data: {
-              status: "PENDING_VERSION",
-              statusReason,
-            },
+    this.$.logger.debug("RunEngine.dequeueFromWorkerQueue(): Pending version", {
+      runId,
+      reason,
+      statusReason,
+    });
+
+    return this.$.runLock.lock("pendingVersion", [runId], 5_000, async (signal) => {
+      this.$.logger.debug("RunEngine.dequeueFromWorkerQueue(): Pending version lock acquired", {
+        runId,
+        reason,
+        statusReason,
+      });
+
+      //mark run as waiting for deploy
+      const run = await prisma.taskRun.update({
+        where: { id: runId },
+        data: {
+          status: "PENDING_VERSION",
+          statusReason,
+        },
+        select: {
+          id: true,
+          status: true,
+          attemptNumber: true,
+          updatedAt: true,
+          createdAt: true,
+          runtimeEnvironment: {
             select: {
               id: true,
-              status: true,
-              attemptNumber: true,
-              updatedAt: true,
-              createdAt: true,
-              runtimeEnvironment: {
-                select: {
-                  id: true,
-                  type: true,
-                  projectId: true,
-                  project: { select: { id: true, organizationId: true } },
-                },
-              },
-            },
-          });
-
-          this.$.logger.debug("RunEngine.dequeueFromWorkerQueue(): Pending version", {
-            runId,
-            run,
-          });
-
-          await this.executionSnapshotSystem.createExecutionSnapshot(prisma, {
-            run,
-            snapshot: {
-              executionStatus: "RUN_CREATED",
-              description:
-                reason ??
-                "The run doesn't have a background worker, so we're going to ack it for now.",
-            },
-            environmentId: run.runtimeEnvironment.id,
-            environmentType: run.runtimeEnvironment.type,
-            projectId: run.runtimeEnvironment.projectId,
-            organizationId: run.runtimeEnvironment.project.organizationId,
-            workerId,
-            runnerId,
-          });
-
-          //we ack because when it's deployed it will be requeued
-          await this.$.runQueue.acknowledgeMessage(orgId, runId);
-
-          this.$.eventBus.emit("runStatusChanged", {
-            time: new Date(),
-            run: {
-              id: runId,
-              status: run.status,
-              updatedAt: run.updatedAt,
-              createdAt: run.createdAt,
-            },
-            organization: {
-              id: run.runtimeEnvironment.project.organizationId,
+              type: true,
+              projectId: true,
+              project: { select: { id: true, organizationId: true } },
             },
-            project: {
-              id: run.runtimeEnvironment.projectId,
-            },
-            environment: {
-              id: run.runtimeEnvironment.id,
-            },
-          });
-        });
-      },
-      {
-        attributes: {
-          runId,
+          },
         },
-      }
-    );
+      });
+
+      this.$.logger.debug("RunEngine.dequeueFromWorkerQueue(): Pending version", {
+        runId,
+        run,
+      });
+
+      await this.executionSnapshotSystem.createExecutionSnapshot(prisma, {
+        run,
+        snapshot: {
+          executionStatus: "RUN_CREATED",
+          description:
+            reason ?? "The run doesn't have a background worker, so we're going to ack it for now.",
+        },
+        environmentId: run.runtimeEnvironment.id,
+        environmentType: run.runtimeEnvironment.type,
+        projectId: run.runtimeEnvironment.projectId,
+        organizationId: run.runtimeEnvironment.project.organizationId,
+        workerId,
+        runnerId,
+      });
+
+      //we ack because when it's deployed it will be requeued
+      await this.$.runQueue.acknowledgeMessage(orgId, runId);
+
+      this.$.eventBus.emit("runStatusChanged", {
+        time: new Date(),
+        run: {
+          id: runId,
+          status: run.status,
+          updatedAt: run.updatedAt,
+          createdAt: run.createdAt,
+        },
+        organization: {
+          id: run.runtimeEnvironment.project.organizationId,
+        },
+        project: {
+          id: run.runtimeEnvironment.projectId,
+        },
+        environment: {
+          id: run.runtimeEnvironment.id,
+        },
+      });
+    });
   }
 }

internal-packages/run-engine/src/engine/tests/heartbeats.test.ts

@@ -350,7 +350,7 @@ describe("RunEngine heartbeats", () => {
         assertNonNullable(executionData2);
         expect(executionData2.snapshot.executionStatus).toBe("QUEUED");
 
-        await setTimeout(1_000);
+        await engine.runQueue.processMasterQueueForEnvironment(authenticatedEnvironment.id);
 
         //have to dequeue again
         const dequeued2 = await engine.dequeueFromWorkerQueue({

internal-packages/run-engine/src/engine/tests/pendingVersion.test.ts

@@ -192,6 +192,7 @@ describe("RunEngine pending version", () => {
         //set this so we have to requeue the runs in two batches
         queueRunsWaitingForWorkerBatchSize: 1,
         tracer: trace.getTracer("test", "0.0.0"),
+        logLevel: "debug",
       });
 
       try {
@@ -247,7 +248,7 @@ describe("RunEngine pending version", () => {
         expect(executionDataR1.snapshot.executionStatus).toBe("QUEUED");
         expect(executionDataR2.snapshot.executionStatus).toBe("QUEUED");
 
-        await setTimeout(500);
+        await engine.runQueue.processMasterQueueForEnvironment(authenticatedEnvironment.id);
 
         //dequeuing should fail
         const dequeued = await engine.dequeueFromWorkerQueue({
@@ -256,6 +257,12 @@ describe("RunEngine pending version", () => {
         });
         expect(dequeued.length).toBe(0);
 
+        const dequeued2 = await engine.dequeueFromWorkerQueue({
+          consumerId: "test_12345",
+          workerQueue: "main",
+        });
+        expect(dequeued2.length).toBe(0);
+
         //queue should be empty
         const queueLength = await engine.runQueue.lengthOfQueue(
           authenticatedEnvironment,