Add more UT and remove repetitive code

Signed-off-by: Bandini Bhopi <bandinib@amazon.com>
opensearch-project · Mar 8, 2024 · 91bf192 · 91bf192
1 parent 53b26ec
commit 91bf192
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 8 deletions.
diff --git a/...plugins/data_source/server/saved_objects/data_source_saved_objects_client_wrapper.test.ts b/...plugins/data_source/server/saved_objects/data_source_saved_objects_client_wrapper.test.ts
@@ -414,12 +414,14 @@ describe('DataSourceSavedObjectsClientWrapper', () => {
     beforeEach(() => {
       mockedClient.update.mockClear();
     });
+
     it('should throw error when pass endpoint to update', async () => {
       const id = 'test1';
       await expect(
         wrapperClient.update(DATA_SOURCE_SAVED_OBJECT_TYPE, id, attributes())
       ).rejects.toThrowError(`Updating a dataSource endpoint is not supported`);
     });
+
     it('should update data source when auth type is present in auth registry', async () => {
       const id = 'test1';
       const mockDataSourceAttributes = attributes({
@@ -448,5 +450,92 @@ describe('DataSourceSavedObjectsClientWrapper', () => {
         expect.anything()
       );
     });
+
+    it('should update throw error when auth type is not present in auth registry', async () => {
+      const id = 'test1';
+      const mockDataSourceAttributes = attributes({
+        auth: {
+          type: 'not_in_registry',
+        },
+      });
+      const { endpoint, ...newObject1 } = mockDataSourceAttributes;
+      mockedClient.get.mockResolvedValue(
+        getSavedObject({
+          id: 'test1',
+          attributes: mockDataSourceAttributes,
+        })
+      );
+      await expect(
+        wrapperClient.update(DATA_SOURCE_SAVED_OBJECT_TYPE, id, newObject1)
+      ).rejects.toThrowError(`Invalid auth type: 'not_in_registry': Bad Request`);
+    });
+  });
+
+  describe('bulkUpdateWithCredentialsEncryption', () => {
+    beforeEach(() => {
+      mockedClient.bulkUpdate.mockClear();
+    });
+
+    it('should update data sources when auth type is present in auth registry', async () => {
+      const mockDataSourceAttributes = attributes({
+        auth: {
+          type: customAuthName,
+        },
+      });
+      const { endpoint, ...bulkUpdateObject } = mockDataSourceAttributes;
+      mockedClient.get.mockResolvedValue(
+        getSavedObject({
+          id: 'test1',
+          attributes: mockDataSourceAttributes,
+        })
+      );
+      await wrapperClient.bulkUpdate(
+        [
+          {
+            id: 'test1',
+            type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+            attributes: bulkUpdateObject,
+          },
+        ],
+        {}
+      );
+      expect(mockedClient.bulkUpdate).toBeCalledWith(
+        [
+          {
+            attributes: bulkUpdateObject,
+            id: 'test1',
+            type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+          },
+        ],
+        {}
+      );
+    });
+
+    it('should bulk update throw error when auth type is not present in auth registry', async () => {
+      const mockDataSourceAttributes = attributes({
+        auth: {
+          type: 'not_in_registry',
+        },
+      });
+      const { endpoint, ...bulkUpdateObject } = mockDataSourceAttributes;
+      mockedClient.get.mockResolvedValue(
+        getSavedObject({
+          id: 'test1',
+          attributes: mockDataSourceAttributes,
+        })
+      );
+      await expect(
+        wrapperClient.bulkUpdate(
+          [
+            {
+              id: 'test1',
+              type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+              attributes: bulkUpdateObject,
+            },
+          ],
+          {}
+        )
+      ).rejects.toThrowError(`Invalid auth type: 'not_in_registry': Bad Request`);
+    });
   });
 });
diff --git a/src/plugins/data_source/server/saved_objects/data_source_saved_objects_client_wrapper.ts b/src/plugins/data_source/server/saved_objects/data_source_saved_objects_client_wrapper.ts
@@ -172,8 +172,7 @@ export class DataSourceSavedObjectsClientWrapper {
           auth: await this.encryptSigV4Credential(auth, { endpoint }),
         };
       default:
-        const authRegistry = (await this.authRegistryPromise).getAuthenticationMethod(auth.type);
-        if (authRegistry !== undefined) {
+        if (await this.isAuthTypeAvailableInRegistry(auth.type)) {
           return attributes;
         }
         throw SavedObjectsErrorHelpers.createBadRequestError(`Invalid auth type: '${auth.type}'`);
@@ -244,8 +243,7 @@ export class DataSourceSavedObjectsClientWrapper {
           return attributes;
         }
       default:
-        const authRegistry = (await this.authRegistryPromise).getAuthenticationMethod(auth.type);
-        if (authRegistry !== undefined) {
+        if (await this.isAuthTypeAvailableInRegistry(auth.type)) {
           return attributes;
         }
         throw SavedObjectsErrorHelpers.createBadRequestError(`Invalid credentials type: '${type}'`);
@@ -338,8 +336,7 @@ export class DataSourceSavedObjectsClientWrapper {
         }
         break;
       default:
-        const authRegistry = (await this.authRegistryPromise).getAuthenticationMethod(type);
-        if (authRegistry !== undefined) {
+        if (await this.isAuthTypeAvailableInRegistry(type)) {
           break;
         }
         throw SavedObjectsErrorHelpers.createBadRequestError(`Invalid auth type: '${type}'`);
@@ -410,8 +407,7 @@ export class DataSourceSavedObjectsClientWrapper {
         encryptionContext = accessKeyEncryptionContext;
         break;
       default:
-        const authRegistry = (await this.authRegistryPromise).getAuthenticationMethod(auth.type);
-        if (authRegistry !== undefined) {
+        if (await this.isAuthTypeAvailableInRegistry(auth.type)) {
           return attributes;
         }
         throw SavedObjectsErrorHelpers.createBadRequestError(`Invalid auth type: '${auth.type}'`);
@@ -494,4 +490,17 @@ export class DataSourceSavedObjectsClientWrapper {
       },
     };
   }
+
+  private async getAuthenticationMethodFromRegistry(type: string) {
+    const authMethod = (await this.authRegistryPromise).getAuthenticationMethod(type);
+    return authMethod;
+  }
+
+  private async isAuthTypeAvailableInRegistry(type: string): Promise<boolean> {
+    const authMethod = await this.getAuthenticationMethodFromRegistry(type);
+    if (authMethod !== undefined) {
+      return true;
+    }
+    return false;
+  }
 }