feat: track IaC local execution tests [CC-972]

Call new endpoint introduced in
https://github.com/snyk/registry/pull/22111 to track billing usage for
IaC tests, which by default run in "local exec" mode.

Author: Craig Furman

src/cli/commands/test/iac-local-execution/analytics.ts

@@ -35,6 +35,7 @@ export enum PerformanceAnalyticsKey {
   OrgSettings = 'org-settings-ms',
   CustomSeverities = 'custom-severities-ms',
   ResultFormatting = 'results-formatting-ms',
+  UsageTracking = 'usage-tracking-ms',
   CacheCleanup = 'cache-cleanup-ms',
   Total = 'total-iac-ms',
 }
@@ -50,6 +51,7 @@ export const performanceAnalyticsObject: Record<
   [PerformanceAnalyticsKey.OrgSettings]: null,
   [PerformanceAnalyticsKey.CustomSeverities]: null,
   [PerformanceAnalyticsKey.ResultFormatting]: null,
+  [PerformanceAnalyticsKey.UsageTracking]: null,
   [PerformanceAnalyticsKey.CacheCleanup]: null,
   [PerformanceAnalyticsKey.Total]: null,
 };

src/cli/commands/test/iac-local-execution/index.ts

@@ -8,6 +8,7 @@ import {
   EngineType,
 } from './types';
 import { addIacAnalytics } from './analytics';
+import { TestLimitReachedError } from './usage-tracking';
 import { TestResult } from '../../../../lib/snyk-test/legacy';
 import {
   initLocalCache,
@@ -17,6 +18,7 @@ import {
   getIacOrgSettings,
   applyCustomSeverities,
   formatScanResults,
+  trackUsage,
   cleanLocalCache,
 } from './measurable-methods';
 import { isFeatureFlagSupportedForOrg } from '../../../../lib/feature-flags';
@@ -62,6 +64,17 @@ export async function test(
       options,
       iacOrgSettings.meta,
     );
+
+    try {
+      await trackUsage(formattedResults);
+    } catch (e) {
+      if (e instanceof TestLimitReachedError) {
+        throw e;
+      }
+      // If something has gone wrong, err on the side of allowing the user to
+      // run their tests by squashing the error.
+    }
+
     addIacAnalytics(formattedResults);
 
     // TODO: add support for proper typing of old TestResult interface.

src/cli/commands/test/iac-local-execution/measurable-methods.ts

@@ -2,6 +2,7 @@ import { loadFiles } from './file-loader';
 import { parseFiles } from './file-parser';
 import { scanFiles } from './file-scanner';
 import { formatScanResults } from './results-formatter';
+import { trackUsage } from './usage-tracking';
 import { cleanLocalCache, initLocalCache } from './local-cache';
 import { applyCustomSeverities } from './org-settings/apply-custom-severities';
 import { getIacOrgSettings } from './org-settings/get-iac-org-settings';
@@ -87,6 +88,11 @@ const measurableFormatScanResults = performanceAnalyticsDecorator(
   PerformanceAnalyticsKey.ResultFormatting,
 );
 
+const measurableTrackUsage = asyncPerformanceAnalyticsDecorator(
+  trackUsage,
+  PerformanceAnalyticsKey.UsageTracking,
+);
+
 const measurableLocalTest = asyncPerformanceAnalyticsDecorator(
   test,
   PerformanceAnalyticsKey.Total,
@@ -100,6 +106,7 @@ export {
   measurableGetIacOrgSettings as getIacOrgSettings,
   measurableApplyCustomSeverities as applyCustomSeverities,
   measurableFormatScanResults as formatScanResults,
+  measurableTrackUsage as trackUsage,
   measurableCleanLocalCache as cleanLocalCache,
   measurableLocalTest as localTest,
 };

src/cli/commands/test/iac-local-execution/usage-tracking.ts

@@ -0,0 +1,55 @@
+import request = require('../../../../lib/request');
+import config = require('../../../../lib/config');
+import { api as getApiToken } from '../../../../lib/api-token';
+import { CustomError } from '../../../../lib/errors';
+
+export async function trackUsage(
+  formattedResults: TrackableResult[],
+): Promise<void> {
+  const trackingData = formattedResults.map((res) => {
+    return {
+      isPrivate: res.meta.isPrivate,
+      issuesPrevented: res.result.cloudConfigResults.length,
+    };
+  });
+  const trackingResponse = await request({
+    method: 'POST',
+    headers: {
+      Authorization: `token ${getApiToken()}`,
+    },
+    url: `${config.API}/track-iac-usage/cli`,
+    body: { results: trackingData },
+    gzip: true,
+    json: true,
+  });
+  switch (trackingResponse.res.statusCode) {
+    case 200:
+      break;
+    case 429:
+      throw new TestLimitReachedError();
+    default:
+      throw new CustomError(
+        'An error occurred while attempting to track test usage: ' +
+          JSON.stringify(trackingResponse.res.body),
+      );
+  }
+}
+
+export class TestLimitReachedError extends CustomError {
+  constructor() {
+    super(
+      'Test limit reached! You have exceeded your infrastructure as code test allocation for this billing period.',
+    );
+  }
+}
+
+// Sub-interface of FormattedResult that we really only use to make test
+// fixtures easier to create.
+export interface TrackableResult {
+  meta: {
+    isPrivate: boolean;
+  };
+  result: {
+    cloudConfigResults: any[];
+  };
+}

test/acceptance/fake-server.ts

@@ -338,6 +338,12 @@ export function fakeServer(root, apikey) {
     return next();
   });
 
+  server.post(root + '/track-iac-usage/cli', (req, res, next) => {
+    res.status(200);
+    res.send({});
+    return next();
+  });
+
   server.setNextResponse = (response) => {
     server._nextResponse = response;
   };

test/jest/unit/iac-unit-tests/usage-tracking.spec.ts

@@ -0,0 +1,95 @@
+import {
+  trackUsage,
+  TestLimitReachedError,
+} from '../../../../src/cli/commands/test/iac-local-execution/usage-tracking';
+import { mocked } from 'ts-jest/utils';
+import { NeedleResponse } from 'needle';
+import makeRequest = require('../../../../src/lib/request/request');
+import { CustomError } from '../../../../src/lib/errors';
+
+jest.mock('../../../../src/lib/request/request');
+const mockedMakeRequest = mocked(makeRequest);
+
+const results = [
+  {
+    meta: {
+      isPrivate: true,
+    },
+    result: {
+      cloudConfigResults: ['an issue'],
+    },
+  },
+  {
+    meta: {
+      isPrivate: false,
+    },
+    result: {
+      cloudConfigResults: [],
+    },
+  },
+];
+
+describe('tracking IaC test usage', () => {
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('does not throw an error when backend returns HTTP 200', async () => {
+    mockedMakeRequest.mockImplementationOnce(() => {
+      return Promise.resolve({
+        res: { statusCode: 200 } as NeedleResponse,
+        body: {
+          foo: 'bar',
+        },
+      });
+    });
+
+    await trackUsage(results);
+
+    expect(mockedMakeRequest.mock.calls.length).toEqual(1);
+    expect(mockedMakeRequest.mock.calls[0][0].body).toEqual({
+      results: [
+        {
+          isPrivate: true,
+          issuesPrevented: 1,
+        },
+        {
+          isPrivate: false,
+          issuesPrevented: 0,
+        },
+      ],
+    });
+  });
+
+  it('throws TestLimitReachedError when backend returns HTTP 429', async () => {
+    mockedMakeRequest.mockImplementationOnce(() => {
+      return Promise.resolve({
+        res: { statusCode: 429 } as NeedleResponse,
+        body: {
+          foo: 'bar',
+        },
+      });
+    });
+
+    await expect(trackUsage(results)).rejects.toThrow(
+      new TestLimitReachedError(),
+    );
+  });
+
+  it('throws CustomError when backend returns HTTP 500', async () => {
+    mockedMakeRequest.mockImplementationOnce(() => {
+      return Promise.resolve({
+        res: { statusCode: 500, body: { foo: 'bar' } } as NeedleResponse,
+        body: {
+          foo: 'bar',
+        },
+      });
+    });
+
+    await expect(trackUsage(results)).rejects.toThrow(
+      new CustomError(
+        'An error occurred while attempting to track test usage: {"foo":"bar"}',
+      ),
+    );
+  });
+});