From 602420e7afefe0350e25dcf282dce6c6c4663d5d Mon Sep 17 00:00:00 2001
From: Warren Parad <gpg@warrenparad.net>
Date: Thu, 29 Aug 2024 09:49:23 +0200
Subject: [PATCH] Support DocumentClient loading from library as well.

---
 src/dynamoDbSafe.js         |  3 ++
 tests/dynamoDbArmor.test.js | 58 +++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/src/dynamoDbSafe.js b/src/dynamoDbSafe.js
index 324b48e..9b85f95 100644
--- a/src/dynamoDbSafe.js
+++ b/src/dynamoDbSafe.js
@@ -157,4 +157,7 @@ class DynamoDB extends DynamoDbOriginal.DocumentClient {
   }
 }
 
+// Allow reference to DynamoDb using the legacy DocumentClient
+DynamoDB.DocumentClient = DynamoDB;
+
 module.exports = { DynamoDB };
diff --git a/tests/dynamoDbArmor.test.js b/tests/dynamoDbArmor.test.js
index 7986009..744cd45 100644
--- a/tests/dynamoDbArmor.test.js
+++ b/tests/dynamoDbArmor.test.js
@@ -40,6 +40,64 @@ describe('dynamoDbArmor.js', () => {
       }
     });
 
+    it('Validate single SET using inner DocumentClient', async () => {
+      const testTable = 'Test-TableId';
+      const testHash = 'testHash';
+      const testRange = 'testRange';
+      const params = {
+        TableName: testTable,
+        Key: {
+          hash: testHash,
+          rang: testRange
+        },
+        UpdateExpression: 'SET #key = :value',
+        ConditionExpression: 'attribute_exists(hash)',
+        ExpressionAttributeNames: {
+          '#key': 'key'
+        },
+        ExpressionAttributeValues: {
+          ':value': 'value'
+        }
+      };
+      try {
+        const dynamoDbOriginalMock = sandbox.mock(DynamoDbOriginal.DocumentClient.prototype);
+        dynamoDbOriginalMock.expects('update').once().returns({ promise() { return Promise.resolve(); } });
+        await new DynamoDB.DocumentClient().update(params);
+        dynamoDbOriginalMock.verify();
+      } catch (error) {
+        expect(error.message).to.eql(null, JSON.stringify(error.message, null, 2));
+      }
+    });
+
+    it('Validate single SET using inner DocumentClient and unnecessary .promise()', async () => {
+      const testTable = 'Test-TableId';
+      const testHash = 'testHash';
+      const testRange = 'testRange';
+      const params = {
+        TableName: testTable,
+        Key: {
+          hash: testHash,
+          rang: testRange
+        },
+        UpdateExpression: 'SET #key = :value',
+        ConditionExpression: 'attribute_exists(hash)',
+        ExpressionAttributeNames: {
+          '#key': 'key'
+        },
+        ExpressionAttributeValues: {
+          ':value': 'value'
+        }
+      };
+      try {
+        const dynamoDbOriginalMock = sandbox.mock(DynamoDbOriginal.DocumentClient.prototype);
+        dynamoDbOriginalMock.expects('update').once().returns({ promise() { return Promise.resolve(); } });
+        await new DynamoDB.DocumentClient().update(params).promise();
+        dynamoDbOriginalMock.verify();
+      } catch (error) {
+        expect(error.message).to.eql(null, JSON.stringify(error.message, null, 2));
+      }
+    });
+
     it('Validate multi SET', async () => {
       const testTable = 'Test-TableId';
       const testHash = 'testHash';