From 2fea84fc369d20eb5eefdd76a3dbff6cc7f4ca33 Mon Sep 17 00:00:00 2001
From: Austin Kelleher <austinlkelleher@gmail.com>
Date: Thu, 21 Apr 2022 07:58:01 -0400
Subject: [PATCH] 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: https://github.com/JupiterOne/graph-qualys/pull/190
---
 .../src/graphObject.ts                        | 15 ++++++++
 .../src/execution/__tests__/jobState.test.ts  | 36 +++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/packages/integration-sdk-private-test-utils/src/graphObject.ts b/packages/integration-sdk-private-test-utils/src/graphObject.ts
index 5575c9e1d..52b65a712 100644
--- a/packages/integration-sdk-private-test-utils/src/graphObject.ts
+++ b/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 {
diff --git a/packages/integration-sdk-runtime/src/execution/__tests__/jobState.test.ts b/packages/integration-sdk-runtime/src/execution/__tests__/jobState.test.ts
index 40e593cfb..c42b2548c 100644
--- a/packages/integration-sdk-runtime/src/execution/__tests__/jobState.test.ts
+++ b/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', () => {