feat: improve error handling & messages when recording to Artillery Cloud (#2531)

hassy · web-flow · commit 339dd3829fe7 · 2024-03-04T17:24:07.000Z
* refactor: enable early shutdowns

Handle cases where we need to shutdown gracefully before a test run starts running.

* feat: improve error checks when recording to Artillery Cloud

* Make error messages more specific
* Check that the API key is valid with a preflight call to cloud API
* Stop early if API key is missing or is invalid
diff --git a/packages/artillery/lib/cmds/run-fargate.js b/packages/artillery/lib/cmds/run-fargate.js
@@ -5,13 +5,12 @@
 const { Command, Flags, Args } = require('@oclif/core');
 const { CommonRunFlags } = require('../cli/common-flags');
 const telemetry = require('../telemetry').init();
-const { Plugin: CloudPlugin } = require('../platform/cloud/cloud');
 
 const runCluster = require('../platform/aws-ecs/legacy/run-cluster');
 const { supportedRegions } = require('../platform/aws-ecs/legacy/util');
 const PlatformECS = require('../platform/aws-ecs/ecs');
 const { ECS_WORKER_ROLE_NAME } = require('../platform/aws/constants');
-
+const { Plugin: CloudPlugin } = require('../platform/cloud/cloud');
 class RunCommand extends Command {
   static aliases = ['run:fargate'];
   // Enable multiple args:
@@ -23,7 +22,34 @@ class RunCommand extends Command {
 
     flags.platform = 'aws:ecs';
 
-    new CloudPlugin(null, null, { flags });
+    const cloud = new CloudPlugin(null, null, { flags });
+    if (cloud.enabled) {
+      try {
+        await cloud.init();
+      } catch (err) {
+        if (err.name === 'CloudAPIKeyMissing') {
+          console.error(
+            'Error: API key is required to record test results to Artillery Cloud'
+          );
+          console.error(
+            'See https://docs.art/get-started-cloud for more information'
+          );
+
+          process.exit(7);
+        } else if (err.name === 'APIKeyUnauthorized') {
+          console.error(
+            'Error: API key is not recognized or is not authorized to record tests'
+          );
+
+          process.exit(7);
+        } else {
+          console.error(
+            'Error: something went wrong connecting to Artillery Cloud'
+          );
+          console.error('Check https://x.com/artilleryio for status updates');
+        }
+      }
+    }
 
     const ECS = new PlatformECS(
       null,
diff --git a/packages/artillery/lib/cmds/run.js b/packages/artillery/lib/cmds/run.js
@@ -111,6 +111,7 @@ RunCommand.args = {
   })
 };
 
+let cloud;
 RunCommand.runCommandImplementation = async function (flags, argv, args) {
   // Collect all input files for reading/parsing - via args, --config, or -i
   const inputFiles = argv.concat(flags.input || [], flags.config || []);
@@ -144,13 +145,43 @@ RunCommand.runCommandImplementation = async function (flags, argv, args) {
   }
 
   try {
+    cloud = new CloudPlugin(null, null, { flags });
+
+    if (cloud.enabled) {
+      try {
+        await cloud.init();
+      } catch (err) {
+        if (err.name === 'CloudAPIKeyMissing') {
+          console.error(
+            'Error: API key is required to record test results to Artillery Cloud'
+          );
+          console.error(
+            'See https://docs.art/get-started-cloud for more information'
+          );
+
+          await gracefulShutdown({ exitCode: 7 });
+        } else if (err.name === 'APIKeyUnauthorized') {
+          console.error(
+            'Error: API key is not recognized or is not authorized to record tests'
+          );
+
+          await gracefulShutdown({ exitCode: 7 });
+        } else {
+          console.error(
+            'Error: something went wrong connecting to Artillery Cloud'
+          );
+          console.error('Check https://x.com/artilleryio for status updates');
+        }
+      }
+    }
+
     const testRunId = process.env.ARTILLERY_TEST_RUN_ID || generateId('t');
     console.log('Test run id:', testRunId);
     global.artillery.testRunId = testRunId;
 
     const script = await prepareTestExecutionPlan(inputFiles, flags, args);
 
-    const runnerOpts = {
+    var runnerOpts = {
       environment: flags.environment,
       // This is used in the worker to resolve
       // the path to the processor module
@@ -197,7 +228,7 @@ RunCommand.runCommandImplementation = async function (flags, argv, args) {
       testRunId
     };
 
-    let launcher = await createLauncher(
+    var launcher = await createLauncher(
       script,
       script.config.payload,
       runnerOpts,
@@ -212,7 +243,7 @@ RunCommand.runCommandImplementation = async function (flags, argv, args) {
       metricsToSuppress
     });
 
-    let reporters = [consoleReporter];
+    var reporters = [consoleReporter];
     if (process.env.CUSTOM_REPORTERS) {
       const customReporterNames = process.env.CUSTOM_REPORTERS.split(',');
       customReporterNames.forEach(function (name) {
@@ -286,8 +317,6 @@ RunCommand.runCommandImplementation = async function (flags, argv, args) {
       }
     });
 
-    new CloudPlugin(null, null, { flags });
-
     global.artillery.globalEvents.emit('test:init', {
       flags,
       testRunId,
@@ -306,8 +335,8 @@ RunCommand.runCommandImplementation = async function (flags, argv, args) {
 
     launcher.run();
 
-    let finalReport = {};
-    let shuttingDown = false;
+    var finalReport = {};
+    var shuttingDown = false;
     process.on('SIGINT', async () => {
       gracefulShutdown({ earlyStop: true });
     });
@@ -353,16 +382,23 @@ RunCommand.runCommandImplementation = async function (flags, argv, args) {
       }
       await Promise.allSettled(ps2);
 
-      await telemetry.shutdown();
+      if (telemetry) {
+        await telemetry.shutdown();
+      }
+
+      if (launcher) {
+        await launcher.shutdown();
+      }
 
-      await launcher.shutdown();
       await (async function () {
-        for (const r of reporters) {
-          if (r.cleanup) {
-            try {
-              await p(r.cleanup.bind(r))();
-            } catch (cleanupErr) {
-              debug(cleanupErr);
+        if (reporters) {
+          for (const r of reporters) {
+            if (r.cleanup) {
+              try {
+                await p(r.cleanup.bind(r))();
+              } catch (cleanupErr) {
+                debug(cleanupErr);
+              }
             }
           }
         }
@@ -565,6 +601,10 @@ async function sendTelemetry(script, flags, extraProps) {
   if (script.config && script.config.__createdByQuickCommand) {
     properties['quick'] = true;
   }
+  if (cloud && cloud.enabled && cloud.user) {
+    properties.cloud = cloud.user;
+  }
+
   properties['solo'] = flags.solo;
   try {
     // One-way hash of target endpoint:
diff --git a/packages/artillery/lib/platform/cloud/cloud.js b/packages/artillery/lib/platform/cloud/cloud.js
@@ -12,22 +12,20 @@ const util = require('node:util');
 
 class ArtilleryCloudPlugin {
   constructor(_script, _events, { flags }) {
+    this.enabled = false;
+
     if (!flags.record) {
       return this;
     }
 
-    this.apiKey = flags.key || process.env.ARTILLERY_CLOUD_API_KEY;
+    this.enabled = true;
 
-    if (!this.apiKey) {
-      console.log(
-        'An API key is required to record test results to Artillery Cloud. See https://docs.art/get-started-cloud for more information.'
-      );
-      return;
-    }
+    this.apiKey = flags.key || process.env.ARTILLERY_CLOUD_API_KEY;
 
     this.baseUrl =
       process.env.ARTILLERY_CLOUD_ENDPOINT || 'https://app.artillery.io';
     this.eventsEndpoint = `${this.baseUrl}/api/events`;
+    this.whoamiEndpoint = `${this.baseUrl}/api/user/whoami`;
 
     this.defaultHeaders = {
       'x-auth-token': this.apiKey
@@ -143,6 +141,10 @@ class ArtilleryCloudPlugin {
     global.artillery.ext({
       ext: 'onShutdown',
       method: async (opts) => {
+        if (!this.enabled || this.off) {
+          return;
+        }
+
         clearInterval(this.setGetLoadTestInterval);
         // Wait for the last logLines events to be processed, as they can sometimes finish processing after shutdown has finished
         await awaitOnEE(
@@ -171,6 +173,44 @@ class ArtilleryCloudPlugin {
     return this;
   }
 
+  async init() {
+    if (!this.apiKey) {
+      const err = new Error();
+      err.name = 'CloudAPIKeyMissing';
+      this.off = true;
+      throw err;
+    }
+
+    let res;
+    let body;
+    try {
+      res = await request.get(this.whoamiEndpoint, {
+        headers: this.defaultHeaders,
+        throwHttpErrors: false,
+        retry: {
+          limit: 0
+        }
+      });
+
+      body = JSON.parse(res.body);
+    } catch (err) {
+      this.off = true;
+      throw err;
+    }
+
+    if (res.statusCode === 401) {
+      const err = new Error();
+      err.name = 'APIKeyUnauthorized';
+      this.off = true;
+      throw err;
+    }
+
+    this.user = {
+      id: body.id,
+      email: body.email
+    };
+  }
+
   async waitOnUnprocessedLogs(maxWaitTime) {
     let waitedTime = 0;
     while (this.unprocessedLogsCounter > 0 && waitedTime < maxWaitTime) {
diff --git a/packages/artillery/test/cli/errors-and-warnings.test.js b/packages/artillery/test/cli/errors-and-warnings.test.js
@@ -55,6 +55,35 @@ tap.test('Suggest similar commands if unknown command is used', async (t) => {
   );
 });
 
+tap.test('Exit early if Artillery Cloud API is not valid', async (t) => {
+  const [exitCode, output] = await execute([
+    'run',
+    '--record',
+    '--key',
+    '123',
+    'test/scripts/gh_215_add_token.json'
+  ]);
+
+  t.equal(exitCode, 7);
+  t.ok(output.stderr.includes('API key is not recognized'));
+});
+
+tap.test(
+  'Exit early if Artillery Cloud API is not valid - on Fargate',
+  async (t) => {
+    const [exitCode, output] = await execute([
+      'run-fargate',
+      '--record',
+      '--key',
+      '123',
+      'test/scripts/gh_215_add_token.json'
+    ]);
+
+    t.equal(exitCode, 7);
+    t.ok(output.stderr.includes('API key is not recognized'));
+  }
+);
+
 /*
  @test "Running a script that uses XPath capture when libxmljs is not installed produces a warning" {
      if [[ ! -z `find . -name "artillery-xml-capture" -type d` ]]; then
