From 2ae64ae481c1e33a5d6db4b90aa8cb9bb9e86941 Mon Sep 17 00:00:00 2001
From: Corey Pyle <corepyle@amazon.com>
Date: Mon, 25 Mar 2024 17:11:55 -0400
Subject: [PATCH] JavaScript (v3): SageMaker - Add afterAll hook in test to
 ensure cleanup on timeouts.

---
 .../SageMakerPipelinesWkflw.js                | 20 +++++++++++--------
 ...er-geospatial-pipeline.integration.test.js | 13 ++++++++++--
 .../example_code/sagemaker/vite.config.js     |  5 +++--
 3 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/javascriptv3/example_code/sagemaker/scenarios/wkflw-sagemaker-geospatial-pipeline/SageMakerPipelinesWkflw.js b/javascriptv3/example_code/sagemaker/scenarios/wkflw-sagemaker-geospatial-pipeline/SageMakerPipelinesWkflw.js
index 574e8dae9fe..a1938649daa 100644
--- a/javascriptv3/example_code/sagemaker/scenarios/wkflw-sagemaker-geospatial-pipeline/SageMakerPipelinesWkflw.js
+++ b/javascriptv3/example_code/sagemaker/scenarios/wkflw-sagemaker-geospatial-pipeline/SageMakerPipelinesWkflw.js
@@ -57,23 +57,27 @@ export class SageMakerPipelinesWkflw {
       console.error(err);
       throw err;
     } finally {
-      // Run all of the clean up functions. If any fail, we log the error and continue.
-      // This ensures all clean up functions are run.
       this.logger.logSeparator();
       const doCleanUp = await this.prompter.confirm({
         message: "Clean up resources?",
       });
       if (doCleanUp) {
-        for (let i = this.cleanUpFunctions.length - 1; i >= 0; i--) {
-          await retry(
-            { intervalInMs: 1000, maxRetries: 60, swallowError: true },
-            this.cleanUpFunctions[i],
-          );
-        }
+        await this.cleanUp();
       }
     }
   }
 
+  async cleanUp() {
+    // Run all of the clean up functions. If any fail, we log the error and continue.
+    // This ensures all clean up functions are run.
+    for (let i = this.cleanUpFunctions.length - 1; i >= 0; i--) {
+      await retry(
+        { intervalInMs: 1000, maxRetries: 60, swallowError: true },
+        this.cleanUpFunctions[i],
+      );
+    }
+  }
+
   async startWorkflow() {
     this.logger.logSeparator(MESSAGES.greetingHeader);
     await this.logger.log(MESSAGES.greeting);
diff --git a/javascriptv3/example_code/sagemaker/tests/wkflw-sagemaker-geospatial-pipeline.integration.test.js b/javascriptv3/example_code/sagemaker/tests/wkflw-sagemaker-geospatial-pipeline.integration.test.js
index 6bace1cf117..ca20105f1d2 100644
--- a/javascriptv3/example_code/sagemaker/tests/wkflw-sagemaker-geospatial-pipeline.integration.test.js
+++ b/javascriptv3/example_code/sagemaker/tests/wkflw-sagemaker-geospatial-pipeline.integration.test.js
@@ -1,7 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import { describe, test } from "vitest";
+import { afterAll, describe, test } from "vitest";
 import { SageMakerPipelinesWkflw } from "../scenarios/wkflw-sagemaker-geospatial-pipeline/SageMakerPipelinesWkflw.js";
 import { Logger } from "@aws-doc-sdk-examples/lib/logger.js";
 import { IAMClient } from "@aws-sdk/client-iam";
@@ -11,6 +11,15 @@ import { LambdaClient } from "@aws-sdk/client-lambda";
 import { SQSClient } from "@aws-sdk/client-sqs";
 
 describe("SageMaker geospatial pipeline workflow", () => {
+  /** @type {SageMakerPipelinesWkflw | undefined} */
+  let wkflw;
+
+  afterAll(async (ctx) => {
+    if (wkflw && ctx.result?.state === "fail") {
+      await wkflw.cleanUp();
+    }
+  });
+
   test("should execute successfully", async () => {
     const prompter = {
       checkContinue() {
@@ -32,7 +41,7 @@ describe("SageMaker geospatial pipeline workflow", () => {
       SQS: new SQSClient({ region: "us-west-2" }),
     };
 
-    const wkflw = new SageMakerPipelinesWkflw(prompter, logger, clients);
+    wkflw = new SageMakerPipelinesWkflw(prompter, logger, clients);
 
     await wkflw.run();
   });
diff --git a/javascriptv3/example_code/sagemaker/vite.config.js b/javascriptv3/example_code/sagemaker/vite.config.js
index f621fdaa735..7c5cef4c90f 100644
--- a/javascriptv3/example_code/sagemaker/vite.config.js
+++ b/javascriptv3/example_code/sagemaker/vite.config.js
@@ -5,11 +5,12 @@ import { defineConfig } from "vitest/config";
 
 const ONE_SECOND = 1000;
 const ONE_MINUTE = ONE_SECOND * 60;
+const ONE_HOUR = ONE_MINUTE * 60;
 
 export default defineConfig({
   test: {
-    testTimeout: ONE_MINUTE * 45,
-    hookTimeout: 50000,
+    testTimeout: ONE_HOUR * 2,
+    hookTimeout: ONE_MINUTE * 30,
     threads: false,
   },
 });