Merge pull request #2522 from artilleryio/feat/async-hooks

artilleryio · Feb 29, 2024 · c2d53ba · c2d53ba
2 parents 095cc1d + 6c8e3f5
commit c2d53ba
Show file tree

Hide file tree

Showing 15 changed files with 179 additions and 68 deletions.
diff --git a/packages/artillery-plugin-memory-inspector/index.js b/packages/artillery-plugin-memory-inspector/index.js
@@ -78,8 +78,6 @@ function ArtilleryPluginMemoryInspector(script, events) {
         continue;
       }
     }
-
-    return next();
   }
 
   script.scenarios = script.scenarios.map((scenario) => {

diff --git a/packages/artillery/lib/platform/aws-ecs/legacy/bom.js b/packages/artillery/lib/platform/aws-ecs/legacy/bom.js
@@ -3,7 +3,7 @@ const fs = require('fs');
 const A = require('async');
 
 const isBuiltinModule = require('is-builtin-module');
-const detective = require('detective');
+const detective = require('detective-es6');
 const depTree = require('dependency-tree');
 
 const walkSync = require('walk-sync');

diff --git a/packages/artillery/lib/platform/local/index.js b/packages/artillery/lib/platform/local/index.js
@@ -103,7 +103,7 @@ class PlatformLocal {
       return {};
     }
 
-    const runnableScript = loadProcessor(
+    const runnableScript = await loadProcessor(
       prepareScript(this.script, _.cloneDeep(this.payload)),
       this.opts
     );

diff --git a/packages/artillery/lib/platform/local/worker.js b/packages/artillery/lib/platform/local/worker.js
@@ -106,7 +106,7 @@ async function prepare(opts) {
   });
 
   const { script: _script, payload, options } = opts;
-  const script = loadProcessor(_script, options);
+  const script = await loadProcessor(_script, options);
 
   global.artillery.testRunId = opts.testRunId;
 

diff --git a/packages/artillery/package.json b/packages/artillery/package.json
@@ -110,7 +110,7 @@
     "csv-parse": "^4.16.3",
     "debug": "^4.3.1",
     "dependency-tree": "^10.0.9",
-    "detective": "^5.1.0",
+    "detective-es6": "^4.0.1",
     "dotenv": "^16.0.1",
     "esbuild-wasm": "^0.19.8",
     "eventemitter3": "^4.0.4",

diff --git a/packages/artillery/test/cli/async-hooks-esm.test.js b/packages/artillery/test/cli/async-hooks-esm.test.js
@@ -0,0 +1,50 @@
+const tap = require('tap');
+const { execute, generateTmpReportPath } = require('../cli/_helpers.js');
+const fs = require('fs');
+
+let reportFilePath;
+tap.beforeEach(async (t) => {
+  reportFilePath = generateTmpReportPath(t.name, 'json');
+});
+
+tap.test('async hooks with ESM', async (t) => {
+  const [exitCode, output] = await execute([
+    'run',
+    '-o',
+    `${reportFilePath}`,
+    'test/scripts/scenario-async-esm-hooks/test.yml'
+  ]);
+
+  t.equal(exitCode, 0, 'CLI should exit with code 0');
+  t.ok(
+    output.stdout.includes('Got context using lodash: true'),
+    'Should be able to use lodash in a scenario to get context'
+  );
+  const json = JSON.parse(fs.readFileSync(reportFilePath, 'utf8'));
+
+  console.log(output);
+
+  t.equal(
+    json.aggregate.counters['http.codes.200'],
+    10,
+    'Should have made 10 requests'
+  );
+
+  t.equal(
+    json.aggregate.counters['hey_from_esm'],
+    10,
+    'Should have emitted 10 custom metrics from ts processor'
+  );
+
+  t.equal(
+    json.aggregate.counters['errors.error_from_async_hook'],
+    10,
+    'Should have emitted 10 errors from an exception in an async hook'
+  );
+
+  t.equal(
+    json.aggregate.counters['vusers.failed'],
+    10,
+    'Should have no completed VUs'
+  );
+});
diff --git a/packages/artillery/test/cli/run-typescript.test.js b/packages/artillery/test/cli/run-typescript.test.js
@@ -95,34 +95,3 @@ tap.test('Runs correctly when package is marked as external', async (t) => {
 
   await deleteFile(bundleLocation);
 });
-
-tap.test(
-  'Failure from a Typescript processor has a resolvable stack trace via source maps',
-  async (t) => {
-    const [exitCode, output] = await execute([
-      'run',
-      '-o',
-      `${reportFilePath}`,
-      'test/scripts/scenarios-typescript/error.yml'
-    ]);
-
-    t.equal(exitCode, 11, 'CLI should exit with code 11');
-    t.ok(
-      output.stdout.includes('error_from_ts_processor'),
-      'Should have logged error from ts processor'
-    );
-
-    // // Search for the path
-    // const pathRegex = /\((.*?):\d+:\d+\)/;
-    // const match = output.stdout.match(pathRegex);
-
-    // // Extract the path if found
-    // const extractedPath = match ? match[1] : null;
-
-    // t.ok(
-    //   extractedPath.includes('.ts'),
-    //   'Should be using source maps to resolve the path to a .ts file'
-    // );
-    // t.ok(fs.existsSync(extractedPath), 'Error path should exist');
-  }
-);
diff --git a/packages/artillery/test/cloud-e2e/fargate/fixtures/ts-external-pkg/processor.ts b/packages/artillery/test/cloud-e2e/fargate/fixtures/ts-external-pkg/processor.ts
@@ -9,13 +9,11 @@ const AddressSchema = z.object({
   country: z.string()
 });
 
-export const checkAddress = async (context, ee, next) => {
+export const checkAddress = async (context, ee) => {
   const address = context.vars.address;
   const result = AddressSchema.safeParse(address);
 
   if (!result.success) {
     ee.emit('error', 'invalid_address');
   }
-
-  next();
 };
diff --git a/packages/artillery/test/cloud-e2e/fargate/run-fargate.test.js b/packages/artillery/test/cloud-e2e/fargate/run-fargate.test.js
@@ -7,6 +7,7 @@ const {
   getTestTags,
   execute
 } = require('../../cli/_helpers.js');
+const path = require('path');
 
 const A9 = process.env.A9 || 'artillery';
 
@@ -224,6 +225,31 @@ test('Run with typescript processor and external package', async (t) => {
   );
 });
 
+test('Run a test with an ESM processor', async (t) => {
+  // The main thing we're checking here is that ESM + dependencies get bundled correctly by BOM
+  const scenarioPath = path.resolve(
+    `${__dirname}/../../scripts/scenario-async-esm-hooks/test.yml`
+  );
+
+  const output =
+    await $`${A9} run-fargate ${scenarioPath} --output ${reportFilePath} --record --tags ${baseTags}`;
+
+  t.equal(output.exitCode, 0, 'CLI exit code should be 0');
+
+  const report = JSON.parse(fs.readFileSync(reportFilePath, 'utf8'));
+  t.equal(
+    report.aggregate.counters['http.codes.200'],
+    10,
+    'Should have made 10 requests'
+  );
+
+  t.equal(
+    report.aggregate.counters['hey_from_esm'],
+    10,
+    'Should have emitted 10 custom metrics from ts processor'
+  );
+});
+
 test('Run lots-of-output', async (t) => {
   $.verbose = false; // we don't want megabytes of output on the console
 

diff --git a/packages/artillery/test/scripts/scenario-async-esm-hooks/helpers.mjs b/packages/artillery/test/scripts/scenario-async-esm-hooks/helpers.mjs
@@ -0,0 +1,11 @@
+import _ from 'lodash';
+
+export const emitCustomMetric = async (context, ee) => {
+  const isESM = _.get(context, 'vars.isESM');
+  console.log(`Got context using lodash: ${JSON.stringify(isESM)}`);
+  ee.emit('counter', 'hey_from_esm', 1);
+};
+
+export const hookThatThrows = async (context, ee) => {
+  throw new Error('error_from_async_hook');
+};
diff --git a/packages/artillery/test/scripts/scenario-async-esm-hooks/test.yml b/packages/artillery/test/scripts/scenario-async-esm-hooks/test.yml
@@ -0,0 +1,16 @@
+config:
+  target: "http://asciiart.artillery.io:8080"
+  phases:
+    - duration: 10
+      arrivalRate: 1
+      name: "Phase 1"
+  processor: "./helpers.mjs"
+  variables:
+    isESM: true
+
+scenarios:
+  - flow:
+      - function: emitCustomMetric
+      - get:
+          url: "/"
+      - function: hookThatThrows
diff --git a/packages/artillery/test/scripts/scenarios-typescript/processor.ts b/packages/artillery/test/scripts/scenarios-typescript/processor.ts
@@ -1,16 +1,13 @@
 import _ from 'lodash';
 
-export const myTest = async (context, ee, next) => {
+export const myTest = async (context, ee) => {
   const isTypescript = _.get(context, 'vars.isTypescript');
 
   console.log(`Got context using lodash: ${JSON.stringify(isTypescript)}`);
 
   ee.emit('counter', 'hey_from_ts', 1);
-
-  next();
 };
 
-export const processorWithError = async (context, ee, next) => {
+export const processorWithError = async (context, ee) => {
   throw new Error('error_from_ts_processor');
-  next();
 };
diff --git a/packages/core/lib/engine_http.js b/packages/core/lib/engine_http.js
@@ -220,9 +220,19 @@ HttpEngine.prototype.step = function step(requestSpec, ee, opts) {
     return function (context, callback) {
       let processFunc = self.config.processor[requestSpec.function];
       if (processFunc) {
-        return processFunc(context, ee, function (hookErr) {
-          return callback(hookErr, context);
-        });
+        if (processFunc.constructor.name === 'Function') {
+          return processFunc(context, ee, function (hookErr) {
+            return callback(hookErr, context);
+          });
+        } else {
+          return processFunc(context, ee)
+            .then(() => {
+              callback(null, context);
+            })
+            .catch((err) => {
+              callback(err, context);
+            });
+        }
       } else {
         debug(`Function "${requestSpec.function}" not defined`);
         debug('processor: %o', self.config.processor);
@@ -309,12 +319,16 @@ HttpEngine.prototype.step = function step(requestSpec, ee, opts) {
           console.log(`WARNING: custom function ${fn} could not be found`); // TODO: a 'warning' event
         }
 
-        processFunc(requestParams, context, ee, function (err) {
-          if (err) {
-            return next(err);
-          }
-          return next(null);
-        });
+        if (processFunc.constructor.name === 'Function') {
+          processFunc(requestParams, context, ee, function (err) {
+            if (err) {
+              return next(err);
+            }
+            return next(null);
+          });
+        } else {
+          processFunc(requestParams, context, ee).then(next).catch(next);
+        }
       },
       function done(err) {
         if (err) {
@@ -608,12 +622,24 @@ HttpEngine.prototype.step = function step(requestSpec, ee, opts) {
                   // Got does not have res.body which Request.js used to have, so we attach it here:
                   res.body = body;
 
-                  processFunc(requestParams, res, context, ee, function (err) {
-                    if (err) {
-                      return next(err);
-                    }
-                    return next(null);
-                  });
+                  if (processFunc.constructor.name === 'Function') {
+                    processFunc(
+                      requestParams,
+                      res,
+                      context,
+                      ee,
+                      function (err) {
+                        if (err) {
+                          return next(err);
+                        }
+                        return next(null);
+                      }
+                    );
+                  } else {
+                    processFunc(requestParams, res, context, ee)
+                      .then(next)
+                      .catch(next);
+                  }
                 },
                 function (err) {
                   if (err) {

diff --git a/packages/core/lib/engine_ws.js b/packages/core/lib/engine_ws.js
@@ -141,9 +141,19 @@ WSEngine.prototype.step = function (requestSpec, ee) {
     return function (context, callback) {
       const processFunc = self.config.processor[requestSpec.function];
       if (processFunc) {
-        processFunc(context, ee, function () {
-          return callback(null, context);
-        });
+        if (processFunc.constructor.name === 'Function') {
+          processFunc(context, ee, function () {
+            return callback(null, context);
+          });
+        } else {
+          return processFunc(context, ee)
+            .then(() => {
+              callback(null, context);
+            })
+            .catch((err) => {
+              callback(err, context);
+            });
+        }
       }
     };
   }

diff --git a/packages/core/lib/runner.js b/packages/core/lib/runner.js
@@ -77,15 +77,25 @@ function loadEngines(
   return { loadedEngines, warnings };
 }
 
-function loadProcessor(script, options) {
+async function loadProcessor(script, options) {
+  const absoluteScriptPath = path.resolve(process.cwd(), options.scriptPath);
   if (script.config.processor) {
-    const absoluteScriptPath = path.resolve(process.cwd(), options.scriptPath);
     const processorPath = path.resolve(
       path.dirname(absoluteScriptPath),
       script.config.processor
     );
-    const processor = require(processorPath);
-    script.config.processor = processor;
+
+    if (processorPath.endsWith('.mjs')) {
+      const exports = await import(processorPath);
+      script.config.processor = Object.assign(
+        {},
+        script.config.processor,
+        exports
+      );
+    } else {
+      // CJS (possibly transplied from TS)
+      script.config.processor = require(processorPath);
+    }
   }
 
   return script;
@@ -439,7 +449,7 @@ function createContext(script, contextVars, additionalProperties = {}) {
         $environment: script._environment,
         $processEnvironment: process.env, // TODO: deprecate
         $env: process.env,
-        $testId: global.artillery.testRunId,
+        $testId: global.artillery.testRunId
       },
       contextVars || {}
     ),
-Original file line number
+Diff line change
@@ Expand Up / @@ -78,8 +78,6 @@ function ArtilleryPluginMemoryInspector(script, events) { @@
             continue;
           }
         }
-        return next();
       }
       script.scenarios = script.scenarios.map((scenario) => {
@@ Expand Down @@