add tests

Signed-off-by: Zhongnan Su <szhongna@amazon.com>

src/plugins/data_source/server/client_pool/client_pool.test.ts

@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { loggingSystemMock } from '../../../../core/server/mocks';
+import { DataSourcePluginConfigType } from '../../config';
+import { OpenSearchClientPool } from './client_pool';
+
+const logger = loggingSystemMock.create();
+
+describe('Client Pool', () => {
+  let service: OpenSearchClientPool;
+  let config: DataSourcePluginConfigType;
+
+  beforeEach(() => {
+    const mockLogger = logger.get('dataSource');
+    service = new OpenSearchClientPool(mockLogger);
+    config = {
+      enabled: true,
+      clientPool: {
+        size: 5,
+      },
+    } as DataSourcePluginConfigType;
+  });
+
+  afterEach(() => {
+    service.stop();
+    jest.clearAllMocks();
+  });
+
+  describe('setup()', () => {
+    test('exposes proper contract', async () => {
+      const setup = await service.setup(config);
+      expect(setup).toHaveProperty('getClientFromPool');
+      expect(setup).toHaveProperty('addClientToPool');
+    });
+  });
+});

src/plugins/data_source/server/client_pool/client_pool.ts

@@ -3,10 +3,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { first } from 'rxjs/operators';
 import { Client } from '@opensearch-project/opensearch';
 import LRUCache from 'lru-cache';
-import { Logger, PluginInitializerContext } from 'src/core/server';
+import { Logger } from 'src/core/server';
 import { DataSourcePluginConfigType } from '../../config';
 
 export interface OpenSearchClientPoolSetup {
@@ -27,15 +26,9 @@ export class OpenSearchClientPool {
   private cache?: LRUCache<string, Client>;
   private isClosed = false;
 
-  constructor(
-    private logger: Logger,
-    private initializerContext: PluginInitializerContext<DataSourcePluginConfigType>
-  ) {}
-
-  public async setup() {
-    const config$ = this.initializerContext.config.create<DataSourcePluginConfigType>();
-    const config: DataSourcePluginConfigType = await config$.pipe(first()).toPromise();
+  constructor(private logger: Logger) {}
 
+  public async setup(config: DataSourcePluginConfigType) {
     const logger = this.logger;
     const { size } = config.clientPool;
 
@@ -78,6 +71,6 @@ export class OpenSearchClientPool {
       return;
     }
     this.isClosed = true;
-    Promise.all(this.cache!.values().map((client) => client.close()));
+    await Promise.all(this.cache!.values().map((client) => client.close()));
   }
 }

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

@@ -0,0 +1,18 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export const ClientMock = jest.fn();
+jest.doMock('@opensearch-project/opensearch', () => {
+  const actual = jest.requireActual('@opensearch-project/opensearch');
+  return {
+    ...actual,
+    Client: ClientMock,
+  };
+});
+
+export const configureClientMock = jest.fn();
+jest.doMock('./configure_client', () => ({
+  configureClient: configureClientMock,
+}));

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

@@ -0,0 +1,101 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import EventEmitter from 'events';
+import { SavedObjectsClientContract } from '../../../core/server';
+import { loggingSystemMock, savedObjectsClientMock } from '../../../core/server/mocks';
+import { CREDENTIAL_SAVED_OBJECT_TYPE, DATA_SOURCE_SAVED_OBJECT_TYPE } from '../common';
+import {
+  CredentialMaterialsType,
+  CredentialSavedObjectAttributes,
+} from '../common/credentials/types';
+import { DataSourceAttributes } from '../common/data_sources';
+import { DataSourcePluginConfigType } from '../config';
+import { OpenSearchClientPoolSetup } from './client_pool/client_pool';
+import { configureClient } from './configure_client';
+import { ClientMock } from './configure_client.test.mocks';
+
+const createFakeClient = () => {
+  const client = new EventEmitter();
+  jest.spyOn(client, 'on');
+  return client;
+};
+
+describe('configureClient', () => {
+  let logger: ReturnType<typeof loggingSystemMock.createLogger>;
+  let config: DataSourcePluginConfigType;
+  let savedObjects: jest.Mocked<SavedObjectsClientContract>;
+  let clientPoolSetup: OpenSearchClientPoolSetup;
+
+  const DATA_SOURCE_ID = 'a54b76ec86771ee865a0f74a305dfff8';
+  const CREDENETIAL_ID = 'a54dsaadasfasfwe22d23d23d2453df3';
+
+  beforeEach(() => {
+    logger = loggingSystemMock.createLogger();
+
+    config = {
+      enabled: true,
+      clientPool: {
+        size: 5,
+      },
+    } as DataSourcePluginConfigType;
+
+    savedObjects = savedObjectsClientMock.create();
+    const dataSourceAttr = {
+      title: 'title',
+      endpoint: 'http://localhost',
+      noAuth: true,
+    } as DataSourceAttributes;
+
+    const crendentialAttr = {
+      title: 'cred',
+      credentialMaterials: {
+        credentialMaterialsType: CredentialMaterialsType.UsernamePasswordType,
+        credentialMaterialsContent: {
+          username: 'username',
+          password: 'password',
+        },
+      },
+    } as CredentialSavedObjectAttributes;
+
+    savedObjects.get.mockResolvedValue({
+      id: DATA_SOURCE_ID,
+      type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+      attributes: dataSourceAttr,
+      references: [],
+    });
+
+    // savedObjects.get.mockResolvedValue({
+    //   id: CREDENETIAL_ID,
+    //   type: CREDENTIAL_SAVED_OBJECT_TYPE,
+    //   attributes: crendentialAttr,
+    //   references: [],
+    // });
+
+    clientPoolSetup = {
+      getClientFromPool: jest.fn(),
+      addClientToPool: jest.fn(),
+    };
+
+    ClientMock.mockImplementation(() => createFakeClient());
+  });
+
+  afterEach(() => {
+    ClientMock.mockReset();
+  });
+
+  test('configure client with noAuth type will call new Client()', async () => {
+    const client = await configureClient(
+      DATA_SOURCE_ID,
+      savedObjects,
+      clientPoolSetup,
+      config,
+      logger
+    );
+
+    expect(ClientMock).toHaveBeenCalledTimes(1);
+    expect(client).toBe(ClientMock.mock.results[0].value);
+  });
+});

src/plugins/data_source/server/configure_client.ts

@@ -0,0 +1,133 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { Client } from '@opensearch-project/opensearch';
+import {
+  Logger,
+  SavedObject,
+  SavedObjectsClientContract,
+  SavedObjectsErrorHelpers,
+} from '../../../../src/core/server';
+import { DATA_SOURCE_SAVED_OBJECT_TYPE, CREDENTIAL_SAVED_OBJECT_TYPE } from '../common';
+import { CredentialSavedObjectAttributes } from '../common/credentials/types';
+import { DataSourceAttributes } from '../common/data_sources';
+import { DataSourcePluginConfigType } from '../config';
+import { OpenSearchClientPoolSetup } from './client_pool/client_pool';
+
+export const configureClient = async (
+  dataSourceId: string,
+  savedObjects: SavedObjectsClientContract,
+  openSearchClientPoolSetup: OpenSearchClientPoolSetup,
+  config: DataSourcePluginConfigType,
+  logger: Logger
+): Promise<Client> => {
+  const dataSource = await getDataSource(dataSourceId, savedObjects);
+  const rootClient = getRootClient(dataSource.attributes, config, openSearchClientPoolSetup);
+
+  return getQueryClient(rootClient, dataSource, savedObjects);
+};
+
+export const getDataSource = async (
+  dataSourceId: string,
+  savedObjects: SavedObjectsClientContract
+): Promise<SavedObject<DataSourceAttributes>> => {
+  try {
+    const dataSource = await savedObjects.get<DataSourceAttributes>(
+      DATA_SOURCE_SAVED_OBJECT_TYPE,
+      dataSourceId
+    );
+    return dataSource;
+  } catch (error: any) {
+    // it will cause 500 error when failed to get saved objects, need to handle such error gracefully
+    throw SavedObjectsErrorHelpers.createBadRequestError(error.message);
+  }
+};
+
+export const getCredential = async (
+  credentialId: string,
+  savedObjects: SavedObjectsClientContract
+): Promise<SavedObject<CredentialSavedObjectAttributes>> => {
+  try {
+    const credential = await savedObjects.get<CredentialSavedObjectAttributes>(
+      CREDENTIAL_SAVED_OBJECT_TYPE,
+      credentialId
+    );
+    return credential;
+  } catch (error: any) {
+    // it will cause 500 error when failed to get saved objects, need to handle such error gracefully
+    throw SavedObjectsErrorHelpers.createBadRequestError(error.message);
+  }
+};
+
+/**
+ * Create a child client object with given auth info.
+ *
+ * @param rootClient root client for the connection with given data source endpoint.
+ * @param dataSource data source saved object
+ * @param savedObjects scoped saved object client
+ * @returns child client.
+ */
+const getQueryClient = async (
+  rootClient: Client,
+  dataSource: SavedObject<DataSourceAttributes>,
+  savedObjects: SavedObjectsClientContract
+): Promise<Client> => {
+  if (dataSource.attributes.noAuth) {
+    return rootClient.child();
+  } else {
+    const credential = await getCredential(dataSource.references[0].id, savedObjects);
+    return getBasicAuthClient(rootClient, credential.attributes);
+  }
+};
+
+/**
+ * Gets a root client object of the OpenSearch endpoint.
+ * Will attempt to get from cache, if cache miss, create a new one and load into cache.
+ *
+ * @param dataSourceAttr data source saved objects attributes.
+ * @returns OpenSearch client for the given data source endpoint.
+ */
+const getRootClient = (
+  dataSourceAttr: DataSourceAttributes,
+  config: DataSourcePluginConfigType,
+  { getClientFromPool, addClientToPool }: OpenSearchClientPoolSetup
+): Client => {
+  const endpoint = dataSourceAttr.endpoint;
+  const cachedClient = getClientFromPool(endpoint);
+  if (cachedClient) {
+    return cachedClient;
+  } else {
+    const client = createClient(config, endpoint);
+
+    addClientToPool(endpoint, client);
+    return client;
+  }
+};
+
+const getBasicAuthClient = (
+  rootClient: Client,
+  credentialAttr: CredentialSavedObjectAttributes
+): Client => {
+  const { username, password } = credentialAttr.credentialMaterials.credentialMaterialsContent;
+  return rootClient.child({
+    auth: {
+      username,
+      password,
+    },
+  });
+};
+
+// TODO: will use client configs, that comes from a merge result of user config and defaults
+export const createClient = (config: DataSourcePluginConfigType, endpoint: string): Client => {
+  const client = new Client({
+    node: endpoint,
+    ssl: {
+      requestCert: true,
+      rejectUnauthorized: true,
+    },
+  });
+
+  return client;
+};

src/plugins/data_source/server/data_source_service.test.ts

@@ -0,0 +1,38 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { loggingSystemMock } from '../../../core/server/mocks';
+import { DataSourcePluginConfigType } from '../config';
+import { DataSourceService } from './data_source_service';
+
+const logger = loggingSystemMock.create();
+
+describe('Data Source Service', () => {
+  let service: DataSourceService;
+  let config: DataSourcePluginConfigType;
+
+  beforeEach(() => {
+    const mockLogger = logger.get('dataSource');
+    service = new DataSourceService(mockLogger);
+    config = {
+      enabled: true,
+      clientPool: {
+        size: 5,
+      },
+    } as DataSourcePluginConfigType;
+  });
+
+  afterEach(() => {
+    service.stop();
+    jest.clearAllMocks();
+  });
+
+  describe('setup()', () => {
+    test('exposes proper contract', async () => {
+      const setup = await service.setup(config);
+      expect(setup).toHaveProperty('getDataSourceClient');
+    });
+  });
+});