From 3bef6642b1304875dab9e3cd764fd9b2700206fa Mon Sep 17 00:00:00 2001
From: Niall Thomson <thomson.niall@gmail.com>
Date: Fri, 9 Aug 2024 10:49:01 -0600
Subject: [PATCH] Fix test utility to inject proper test output to after hooks

---
 test/util/src/lib/markdownsh.ts               | 23 +++++++++--------
 test/util/src/lib/shell/shell.ts              | 25 +++++++++++++------
 test/util/test-content/basics/expect-error.md |  6 +++++
 .../basics/tests/hook-expect-error.sh         | 16 ++++++++++++
 4 files changed, 53 insertions(+), 17 deletions(-)
 create mode 100644 test/util/test-content/basics/tests/hook-expect-error.sh

diff --git a/test/util/src/lib/markdownsh.ts b/test/util/src/lib/markdownsh.ts
index d6a84c2e6..8b3d50a9e 100644
--- a/test/util/src/lib/markdownsh.ts
+++ b/test/util/src/lib/markdownsh.ts
@@ -294,9 +294,15 @@ class CustomTest extends Test {
                 {},
               );
             } catch (e: any) {
-              e.message = `Error running test case command at line ${testCase.lineNumber} - ${e.message}`;
-
-              throw e;
+              if (e instanceof ShellError && testCase.expectError) {
+                if (debug) {
+                  console.log("Ignoring expected error");
+                }
+              } else {
+                e.message = `Error running test case command at line ${testCase.lineNumber} - ${e.message}`;
+
+                throw e;
+              }
             }
 
             await this.hook(testCase, category, "after", hookTimeout, {
@@ -316,13 +322,10 @@ class CustomTest extends Test {
                 `Script failed to complete within ${e.timeout} seconds`,
               );
             } else if (e instanceof ShellError) {
-              if (!testCase.expectError) {
-                console.log(e.message);
-                console.log(`Command returned error code ${e.code}`);
-                console.log(`stdout: \n${e.stdout}`);
-                console.log(`stderr: \n${e.stderr}`);
-                assert.fail("Script exit with an error code");
-              }
+              console.log(e.message);
+              console.log(`Command returned error code ${e.code}`);
+              console.log(`output: \n${e.output}`);
+              assert.fail("Script exit with an error code");
             } else {
               assert.fail(`An unknown error occurred: ${e.message}`);
             }
diff --git a/test/util/src/lib/shell/shell.ts b/test/util/src/lib/shell/shell.ts
index 3e0eef018..05cbf2c1d 100644
--- a/test/util/src/lib/shell/shell.ts
+++ b/test/util/src/lib/shell/shell.ts
@@ -57,8 +57,6 @@ export class DefaultShell implements Shell {
       let processingFunction = false;
       let env: { [key: string]: string } = {};
 
-      let processedOutput = "";
-
       for (let step = 0; step < parts.length; step++) {
         const line = parts[step];
 
@@ -80,12 +78,18 @@ export class DefaultShell implements Shell {
         } else {
           if (line === DefaultShell.ENV_MARKER) {
             processingEnv = true;
-          } else {
-            processedOutput += `${line}\n`;
           }
         }
       }
 
+      let finalOutput = "";
+
+      let outputParts = buffer.toString().split(DefaultShell.ENV_MARKER);
+
+      if (outputParts.length > 1) {
+        finalOutput = outputParts[0];
+      }
+
       if (processingEnv) {
         this.environment = {
           ...this.environment,
@@ -93,7 +97,7 @@ export class DefaultShell implements Shell {
         };
       }
 
-      return Promise.resolve(new ExecutionResult(processedOutput));
+      return Promise.resolve(new ExecutionResult(finalOutput));
     } catch (e: any) {
       if (e.code) {
         throw new ShellTimeout(
@@ -105,10 +109,16 @@ export class DefaultShell implements Shell {
       }
 
       if (!expect) {
-        throw new ShellError(e.status, e.message, e.stdout, e.stderr);
+        throw new ShellError(
+          e.status,
+          `Command failed: ${command}`,
+          e.stdout,
+          e.stderr,
+          e.output,
+        );
       }
 
-      return Promise.resolve(new ExecutionResult(e.stderr));
+      return Promise.resolve(new ExecutionResult(e.output));
     }
   }
 }
@@ -119,6 +129,7 @@ export class ShellError extends Error {
     message: string,
     public stdout: string,
     public stderr: string,
+    public output: string,
   ) {
     super(message);
 
diff --git a/test/util/test-content/basics/expect-error.md b/test/util/test-content/basics/expect-error.md
index 2c08ded9c..6c929f433 100644
--- a/test/util/test-content/basics/expect-error.md
+++ b/test/util/test-content/basics/expect-error.md
@@ -9,3 +9,9 @@ Scenario:
 $ eklasmd0ajs0dasipod
 # Under normal circumstances this would fail
 ```
+
+```bash expectError=true hook=expect-error
+$ ls -la
+$ eklasmd0ajs0dasipod
+# Under normal circumstances this would fail
+```
diff --git a/test/util/test-content/basics/tests/hook-expect-error.sh b/test/util/test-content/basics/tests/hook-expect-error.sh
new file mode 100644
index 000000000..9ccffa2a7
--- /dev/null
+++ b/test/util/test-content/basics/tests/hook-expect-error.sh
@@ -0,0 +1,16 @@
+set -e
+
+before() {
+  echo "This hook executes before the test"
+}
+
+after() {
+  if [[ $TEST_OUTPUT != *"command not found"* ]]; then
+    echo "Failed to match expected output"
+    echo $TEST_OUTPUT
+
+    exit 1
+  fi
+}
+
+"$@"