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

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-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 await 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', () => {