Add UT for configure client with Auth registry

Signed-off-by: Bandini Bhopi <bandinib@amazon.com>

src/core/server/data_source/authentication_methods_registry.mock.ts

@@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { IAuthenticationMethodRegistery } from './authentication_methods_registry';
+
+const create = () =>
+  (({
+    getAllAuthenticationMethods: jest.fn(),
+    getAuthenticationMethod: jest.fn(),
+  } as unknown) as jest.Mocked<IAuthenticationMethodRegistery>);
+
+export const authenticationMethodRegisteryMock = { create };

src/core/server/index.ts

@@ -83,6 +83,7 @@ import {
   DataSourceServiceStart,
   AuthenticationMethod,
   IAuthenticationMethodRegistery,
+  AuthAttributes,
 } from './data_source';
 
 // Because of #79265 we need to explicity import, then export these types for
@@ -511,6 +512,7 @@ export {
   CrossCompatibilityServiceStart,
   AuthenticationMethod,
   IAuthenticationMethodRegistery,
+  AuthAttributes,
 };
 
 /**

src/core/server/mocks.ts

@@ -73,6 +73,7 @@ export { contextServiceMock } from './context/context_service.mock';
 export { capabilitiesServiceMock } from './capabilities/capabilities_service.mock';
 export { coreUsageDataServiceMock } from './core_usage_data/core_usage_data_service.mock';
 export { crossCompatibilityServiceMock } from './cross_compatibility/cross_compatibility.mock';
+export { authenticationMethodRegisteryMock } from './data_source/authentication_methods_registry.mock';
 
 export function pluginInitializerContextConfigMock<T>(config: T) {
   const globalConfig: SharedGlobalConfig = {

src/plugins/data_source/server/client/configure_client.test.ts

@@ -3,8 +3,16 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { SavedObjectsClientContract } from '../../../../core/server';
-import { loggingSystemMock, savedObjectsClientMock } from '../../../../core/server/mocks';
+import {
+  SavedObjectsClientContract,
+  IAuthenticationMethodRegistery,
+  AuthenticationMethod,
+} from '../../../../core/server';
+import {
+  loggingSystemMock,
+  savedObjectsClientMock,
+  authenticationMethodRegisteryMock,
+} from '../../../../core/server/mocks';
 import { DATA_SOURCE_SAVED_OBJECT_TYPE } from '../../common';
 import {
   DataSourceAttributes,
@@ -38,12 +46,14 @@ describe('configureClient', () => {
   let dataSourceClientParams: DataSourceClientParams;
   let usernamePasswordAuthContent: UsernamePasswordTypedContent;
   let sigV4AuthContent: SigV4Content;
+  let authenticationMethodRegistery: jest.Mocked<IAuthenticationMethodRegistery>;
 
   beforeEach(() => {
     dsClient = opensearchClientMock.createInternalClient();
     logger = loggingSystemMock.createLogger();
     savedObjectsMock = savedObjectsClientMock.create();
     cryptographyMock = cryptographyServiceSetupMock.create();
+    authenticationMethodRegistery = authenticationMethodRegisteryMock.create();
 
     config = {
       enabled: true,
@@ -238,4 +248,54 @@ describe('configureClient', () => {
     expect(savedObjectsMock.get).toHaveBeenCalledTimes(1);
     expect(decodeAndDecryptSpy).toHaveBeenCalledTimes(1);
   });
+
+  test('should throw error when no credentials present for auth method other than NO_AUTH', async () => {
+    savedObjectsMock.get.mockReset().mockResolvedValueOnce({
+      id: DATA_SOURCE_ID,
+      type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+      attributes: {
+        ...dataSourceAttr,
+        auth: {
+          type: 'token_exchange',
+        },
+      },
+      references: [],
+    });
+    await expect(
+      configureClient(dataSourceClientParams, clientPoolSetup, config, logger)
+    ).rejects.toThrowError();
+  });
+
+  test('configureClient should retunrn client from authentication registery if method present in registry', async () => {
+    const tokenExchangeContent = {
+      region: 'us-east-1',
+      roleARN: 'test-role',
+    };
+    savedObjectsMock.get.mockReset().mockResolvedValueOnce({
+      id: DATA_SOURCE_ID,
+      type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+      attributes: {
+        ...dataSourceAttr,
+        auth: {
+          type: 'token_exchange',
+          credentials: tokenExchangeContent,
+        },
+      },
+      references: [],
+    });
+    const authMethod: AuthenticationMethod = { authType: 'token_exchange', getClient: ClientMock };
+    authenticationMethodRegistery.getAuthenticationMethod.mockImplementation(() => authMethod);
+
+    const client = await configureClient(
+      dataSourceClientParams,
+      clientPoolSetup,
+      config,
+      logger,
+      authenticationMethodRegistery
+    );
+    expect(authenticationMethodRegistery.getAuthenticationMethod).toHaveBeenCalledTimes(1);
+    expect(ClientMock).toHaveBeenCalledTimes(1);
+    expect(savedObjectsMock.get).toHaveBeenCalledTimes(1);
+    // expect(client).toBe(dsClient.child.mock.results[0].value);
+  });
 });

src/plugins/data_source/server/client/configure_client.ts

@@ -44,7 +44,7 @@ export const configureClient = async (
   openSearchClientPoolSetup: OpenSearchClientPoolSetup,
   config: DataSourcePluginConfigType,
   logger: Logger,
-  registeredAuthRegistry: IAuthenticationMethodRegistery | undefined
+  registeredAuthRegistry?: IAuthenticationMethodRegistery | undefined
 ): Promise<Client> => {
   let dataSource;
   let requireDecryption = true;
@@ -126,7 +126,7 @@ const getQueryClient = async (
   const clientOptions = parseClientOptions(config, endpoint);
   const cacheKey = generateCacheKey(dataSourceAttr, dataSourceId);
 
-  if (!dataSourceAttr.auth.credentials) {
+  if (type !== AuthType.NoAuth && !dataSourceAttr.auth.credentials) {
     throw Error(`Credentials not found.`);
   }
   if (registeredAuthRegistry !== undefined) {