From 59cc2144ddae2449a7e82a6b7e61edbdca03b550 Mon Sep 17 00:00:00 2001
From: Anthony Heber <anthony@heber.dev>
Date: Thu, 1 Aug 2024 08:59:50 -0600
Subject: [PATCH] v0.0.15 improve timeout handling

---
 package.json             |  2 +-
 src/lib/commands/apex.ts |  8 +++++++-
 src/lib/polling.ts       | 17 ++++++++++++-----
 src/lib/sf.ts            | 22 +++++++++++++++++++---
 4 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/package.json b/package.json
index 123a373..c3222d8 100644
--- a/package.json
+++ b/package.json
@@ -1,7 +1,7 @@
 {
   "name": "sf-warp",
   "description": "Warp your code and see if your tests notice",
-  "version": "0.0.14",
+  "version": "0.0.15",
   "dependencies": {
     "@oclif/core": "^3.15.1",
     "@salesforce/core": "^6.4.2",
diff --git a/src/lib/commands/apex.ts b/src/lib/commands/apex.ts
index 1c1ae46..cf41f1e 100644
--- a/src/lib/commands/apex.ts
+++ b/src/lib/commands/apex.ts
@@ -231,7 +231,13 @@ export default class ApexWarper {
 
           if (error instanceof Error) {
             errorMessage = error.message;
-            if (this.atLeastVerbosity(Verbosity.minimal)) {
+            if (error.message === 'Timeout polling action') {
+              if (this.config.verbosity === Verbosity.minimal) {
+                process.stdout.write('⏰');
+              } else if (this.atLeastVerbosity(Verbosity.details)) {
+                console.log('Test Timed Out ⏰');
+              }
+            } else if (this.atLeastVerbosity(Verbosity.minimal)) {
               console.log('Failure:', error.message || error);
             }
           }
diff --git a/src/lib/polling.ts b/src/lib/polling.ts
index 259829b..604b1b6 100644
--- a/src/lib/polling.ts
+++ b/src/lib/polling.ts
@@ -3,7 +3,7 @@ interface PollConfig {
   timeout?: number;
   actionName: string;
   action(): unknown;
-  cancelAction?(): void;
+  cancelAction?(): void | Promise<void>;
 }
 
 const cancelledTimeouts: NodeJS.Timeout[] = [];
@@ -13,7 +13,17 @@ export function pollForResult<T>(pollConfig: PollConfig): Promise<T> {
   let timeoutId: NodeJS.Timeout;
   return new Promise((resolve, reject): void => {
     timeoutId = setTimeout(() => {
-      reject('Timeout polling action');
+      let cancelPromise: Promise<void> | undefined;
+      if (pollConfig.cancelAction) {
+        cancelPromise = (pollConfig.cancelAction.call(null) as Promise<void>).catch(reject);
+      }
+      if (cancelPromise) {
+        void cancelPromise.finally(() => {
+          reject(new Error('Timeout polling action'));
+        });
+      } else {
+        reject(new Error('Timeout polling action'));
+      }
       cancelledTimeouts.push(timeoutId);
     }, pollConfig.timeout ?? 30000);
     void executePollAction<T>(pollConfig, resolve, reject, timeoutId);
@@ -29,9 +39,6 @@ function executePollAction<T>(
   pollTime?: number,
 ): void {
   if (cancelledTimeouts.includes(timeoutId)) {
-    if (pollConfig.cancelAction) {
-      pollConfig.cancelAction.call(null);
-    }
     return;
   }
   const waitTime = pollTime ?? pollConfig.initialWaitMs ?? 1000;
diff --git a/src/lib/sf.ts b/src/lib/sf.ts
index e923c9c..d3cd434 100644
--- a/src/lib/sf.ts
+++ b/src/lib/sf.ts
@@ -1,6 +1,5 @@
 import { Connection } from '@salesforce/core';
-// eslint-disable-next-line import/no-extraneous-dependencies
-import { Record } from 'jsforce';
+import type { Record, Schema, SObjectUpdateRecord } from 'jsforce';
 import { pollForResult } from './polling';
 
 interface ApexTestRunResultRecord extends Record {
@@ -13,6 +12,10 @@ interface ApexTestRunResultRecord extends Record {
   MethodsFailed: number;
 }
 
+interface ApexTestQueueItem extends Record {
+  Status: string;
+}
+
 interface ContainerAsyncRequestRecord extends Record {
   State: string;
   ErrorMsg: string;
@@ -37,7 +40,7 @@ export async function executeTests(
   testClasses: string[],
   timeoutMs: number,
 ): Promise<ApexTestRunResultRecord> {
-  const asyncJobId = await conn.tooling.runTestsAsynchronous({ classNames: testClasses.join(',') });
+  const asyncJobId = await conn.tooling.runTestsAsynchronous({ classNames: testClasses.join(','), maxFailedTests: 0 });
 
   return pollForResult({
     timeout: timeoutMs,
@@ -50,6 +53,19 @@ export async function executeTests(
         return request.records[0];
       }
     },
+    cancelAction: async () => {
+      const queryResults = await conn.query<ApexTestQueueItem>(
+        `SELECT Id, Status FROM ApexTestQueueItem WHERE ParentJobId = '${asyncJobId}'`,
+      );
+      queryResults.records.forEach((record) => (record.Status = 'Aborted'));
+      const saveResult = await conn.update(
+        'ApexTestQueueItem',
+        queryResults.records as Array<SObjectUpdateRecord<Schema, 'ApexTestQueueItem'>>,
+      );
+      if (saveResult.find((res) => res.success === false) === undefined) {
+        throw new Error('Failure to cancel tests:' + JSON.stringify(saveResult));
+      }
+    },
   });
 }