INT-3404 - Add tests to validate behavior of jobState.hasKey

The Qualys integration was using `await`ing the result of `jobState.hasKey`
and firing off multiple promises with [`pQueue`](https://github.com/sindresorhus/p-queue).
The Qualys integration was occasionally seeing `DUPLICATE_KEY_DETECTED`
errors being thrown despite checking whether the graph object `_key`
exists before adding the new entity to the `jobState`. These tests were
added as part of the investigation into the issue.

The Qualys integration fix can be found here: JupiterOne-Archives/graph-qualys#190

packages/integration-sdk-private-test-utils/src/graphObject.ts

@@ -1,4 +1,5 @@
 import { Entity, ExplicitRelationship } from '@jupiterone/integration-sdk-core';
+import { times } from 'lodash';
 import { v4 as uuid } from 'uuid';
 
 export function createTestEntity(partial?: Partial<Entity>): Entity {
@@ -11,6 +12,20 @@ export function createTestEntity(partial?: Partial<Entity>): Entity {
   };
 }
 
+/**
+ * Create `n` test entities
+ *
+ * @param n {number} Number of test entities to create
+ * @returns {Entity[]}
+ */
+export function createTestEntities(n: number): Entity[] {
+  return times(n, (i) =>
+    createTestEntity({
+      _key: `entityKey:${i}`,
+    }),
+  );
+}
+
 export function createTestRelationship(
   partial?: Partial<ExplicitRelationship>,
 ): ExplicitRelationship {

packages/integration-sdk-private-test-utils/src/util.ts

@@ -3,3 +3,20 @@ export function sleep(ms: number) {
     setTimeout(() => resolve(undefined), ms);
   });
 }
+
+/**
+ * Utility to invoke a callback `n` times. Replaces need to install `lodash.times`
+ *
+ * @param n {number} Number of times to run the callback function
+ * @param cb {Function} Callback function to invoke `n` times
+ * @returns
+ */
+export function times<T>(n: number, cb: (i: number) => T) {
+  const all: T[] = [];
+
+  for (let i = 0; i < n; i++) {
+    all.push(cb(i));
+  }
+
+  return all;
+}

packages/integration-sdk-runtime/src/execution/__tests__/jobState.test.ts

@@ -14,12 +14,14 @@ import {
   createTestEntity,
   createTestRelationship,
   sleep,
+  createTestEntities,
 } from '@jupiterone/integration-sdk-private-test-utils';
 import {
   createQueuedStepGraphObjectDataUploader,
   CreateQueuedStepGraphObjectDataUploaderParams,
 } from '../uploader';
 import { FlushedGraphObjectData } from '../../storage/types';
+import pMap from 'p-map';
 
 jest.mock('fs');
 
@@ -172,6 +174,40 @@ describe('#hasKey', () => {
     expect(await jobState.hasKey('A')).toBeTrue();
     expect(await jobState.hasKey('a')).toBeTrue();
   });
+
+  test('should handle concurrent reads from in-memory key store when using synchronous hasKey', async () => {
+    const jobState = createTestStepJobState();
+    const entities = createTestEntities(100);
+
+    const tenthEntity = entities[9];
+    entities.splice(10, 0, tenthEntity);
+
+    const results = await pMap(entities, async (e) => {
+      if (jobState.hasKey(e._key)) return;
+      return jobState.addEntity(e);
+    });
+
+    expect(results.length).toEqual(entities.length);
+  });
+
+  test('should fail to handle concurrent reads from in-memory key store when using asynchronous hasKey', async () => {
+    const jobState = createTestStepJobState();
+    const entities = createTestEntities(100);
+
+    const tenthEntity = entities[9];
+    entities.splice(10, 0, tenthEntity);
+
+    await expect(
+      pMap(entities, async (e) => {
+        // NOTE: Due to the event loop queue order in the Node.js, awaiting
+        // the `hasKey` promise while handling multiple entities concurrently,
+        // could result in a `hasKey` returning `false` because neither of the
+        // duplicate entities have been fully added to the job state yet.
+        if (await jobState.hasKey(e._key)) return;
+        await jobState.addEntity(e);
+      }),
+    ).rejects.toThrowError('Duplicate _key detected (_key=entityKey:9)');
+  });
 });
 
 describe('upload callbacks', () => {