Skip to content

Commit

Permalink
Added UT for client pool
Browse files Browse the repository at this point in the history 
Signed-off-by: Bandini Bhopi <bandinib@amazon.com>
  • Loading branch information
@bandinib-amzn
bandinib-amzn committed Mar 14, 2024
1 parent 2c2e7cf commit 9933938
Show file tree
Hide file tree
Showing 8 changed files with 758 additions and 34 deletions.
8 changes: 8 additions & 0 deletions src/plugins/data_source/server/auth_registry/authentication_methods_registry.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,21 @@
import { AuthenticationMethodRegistery } from './authentication_methods_registry';
import { AuthenticationMethod } from '../../server/types';
import { AuthType } from '../../common/data_sources';
import { OpenSearchClientPoolSetup } from '../client';

const clientPoolSetup: OpenSearchClientPoolSetup = {
getClientFromPool: jest.fn(),
addClientToPool: jest.fn(),
};

const createAuthenticationMethod = (
authMethod: Partial<AuthenticationMethod>
): AuthenticationMethod => ({
name: 'unknown',
authType: AuthType.NoAuth,
credentialProvider: jest.fn(),
clientPoolSetup,
legacyClientPoolSetup: clientPoolSetup,
...authMethod,
});

Expand Down
10 changes: 10 additions & 0 deletions src/plugins/data_source/server/auth_registry/authentication_methods_registry.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import { deepFreeze } from '@osd/std';
import { AuthenticationMethod } from '../../server/types';
import { AuthType } from '../../common/data_sources';

export type IAuthenticationMethodRegistery = Omit<
AuthenticationMethodRegistery,
Expand All @@ -18,6 +19,15 @@ export class AuthenticationMethodRegistery {
* Authentication Method can only be registered once. subsequent calls with the same method name will throw an error.
*/
public registerAuthenticationMethod(method: AuthenticationMethod) {
if (
method.name === AuthType.NoAuth ||
method.name === AuthType.UsernamePasswordType ||
method.name === AuthType.SigV4
) {
throw new Error(

Check warning on line 27 in src/plugins/data_source/server/auth_registry/authentication_methods_registry.ts

View check run for this annotation

Codecov / codecov/patch

src/plugins/data_source/server/auth_registry/authentication_methods_registry.ts#L27 

Added line #L27 was not covered by tests
`Must not be no_auth or username_password or sigv4 for registered auth types`
);
}
if (this.authMethods.has(method.name)) {
throw new Error(`Authentication method '${method.name}' is already registered`);
}
Expand Down
6 changes: 5 additions & 1 deletion src/plugins/data_source/server/client/client_pool.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,11 @@ export class OpenSearchClientPool {
});
this.logger.info(`Created data source aws client pool of size ${size}`);

const getClientFromPool = (key: string, authType: AuthType) => {
const getClientFromPool = (
key: string,
authType: AuthType,
request?: OpenSearchDashboardsRequest
) => {
const selectedCache = authType === AuthType.SigV4 ? this.awsClientCache : this.clientCache;

return selectedCache!.get(key);
Expand Down
278 changes: 277 additions & 1 deletion src/plugins/data_source/server/client/configure_client.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ import {
parseClientOptionsMock,
authRegistryCredentialProviderMock,
} from './configure_client.test.mocks';
import { OpenSearchClientPoolSetup } from './client_pool';
import { OpenSearchClientPool, OpenSearchClientPoolSetup } from './client_pool';
import { configureClient } from './configure_client';
import { ClientOptions } from '@opensearch-project/opensearch';
// eslint-disable-next-line @osd/eslint/no-restricted-paths
Expand Down Expand Up @@ -361,4 +361,280 @@ describe('configureClient', () => {
},
});
});

test('configure client with auth method from registry, service == aoss, should successfully call new Client()', async () => {
savedObjectsMock.get.mockReset().mockResolvedValueOnce({
id: DATA_SOURCE_ID,
type: DATA_SOURCE_SAVED_OBJECT_TYPE,
attributes: {
...dataSourceAttr,
auth: {
type: AuthType.SigV4,
credentials: { ...customAuthContent, service: 'aoss' },
},
},
references: [],
});

authRegistryCredentialProviderMock.mockReturnValue({
credential: sigV4AuthContent,
type: AuthType.SigV4,
});

const client = await configureClient(
{ ...dataSourceClientParams, authRegistry: authenticationMethodRegistery },
clientPoolSetup,
config,
logger
);

expect(ClientMock).toHaveBeenCalledTimes(1);
expect(client).toBe(dsClient.child.mock.results[0].value);
expect(dsClient.child).toBeCalledWith({
auth: {
credentials: {
accessKeyId: sigV4AuthContent.accessKey,
secretAccessKey: sigV4AuthContent.secretKey,
sessionToken: '',
},
region: sigV4AuthContent.region,
service: 'aoss',
},
});
});

describe('Client Pool', () => {
let opensearchClientPoolSetup: OpenSearchClientPoolSetup;
let openSearchClientPool: OpenSearchClientPool;
beforeEach(() => {
openSearchClientPool = new OpenSearchClientPool(logger);
opensearchClientPoolSetup = openSearchClientPool.setup(config);
});

describe('NoAuth', () => {
beforeEach(() => {
savedObjectsMock.get.mockReset().mockResolvedValue({
id: DATA_SOURCE_ID,
type: DATA_SOURCE_SAVED_OBJECT_TYPE,
attributes: {
...dataSourceAttr,
auth: {
type: AuthType.NoAuth,
},
},
references: [],
});
});

test('For same endpoint only one client object should be created', async () => {
await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);
await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);

expect(ClientMock).toHaveBeenCalledTimes(1);
});

test('For different endpoints multiple client objects should be created', async () => {
await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);

const mockDataSourceAttr = { ...dataSourceAttr, endpoint: 'http://test.com' };

savedObjectsMock.get.mockReset().mockResolvedValueOnce({
id: DATA_SOURCE_ID,
type: DATA_SOURCE_SAVED_OBJECT_TYPE,
attributes: {
...mockDataSourceAttr,
auth: {
type: AuthType.NoAuth,
},
},
references: [],
});

await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);

expect(ClientMock).toHaveBeenCalledTimes(2);
});
});

describe('UserNamePassword', () => {
beforeEach(() => {
savedObjectsMock.get.mockReset().mockResolvedValue({
id: DATA_SOURCE_ID,
type: DATA_SOURCE_SAVED_OBJECT_TYPE,
attributes: dataSourceAttr,
references: [],
});
jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({
decryptedText: 'password',
encryptionContext: { endpoint: 'http://localhost' },
});
});

test('For same endpoint only one client object should be created', async () => {
await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);
await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);

expect(ClientMock).toHaveBeenCalledTimes(1);
});

test('For different endpoints multiple client objects should be created', async () => {
await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);

const mockDataSourceAttr = { ...dataSourceAttr, endpoint: 'http://test.com' };
savedObjectsMock.get.mockReset().mockResolvedValue({
id: DATA_SOURCE_ID,
type: DATA_SOURCE_SAVED_OBJECT_TYPE,
attributes: mockDataSourceAttr,
references: [],
});
jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({
decryptedText: 'password',
encryptionContext: { endpoint: 'http://test.com' },
});

await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);

expect(ClientMock).toHaveBeenCalledTimes(2);
});
});

describe('AWSSigV4', () => {
beforeEach(() => {
savedObjectsMock.get.mockReset().mockResolvedValue({
id: DATA_SOURCE_ID,
type: DATA_SOURCE_SAVED_OBJECT_TYPE,
attributes: {
...dataSourceAttr,
auth: {
type: AuthType.SigV4,
credentials: sigV4AuthContent,
},
},
references: [],
});

jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({
decryptedText: 'accessKey',
encryptionContext: { endpoint: 'http://localhost' },
});
});
test('For same endpoint only one client object should be created', async () => {
await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);
await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);

expect(ClientMock).toHaveBeenCalledTimes(1);
});

test('For different endpoints multiple client objects should be created', async () => {
await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);

const mockDataSourceAttr = { ...dataSourceAttr, endpoint: 'http://test.com' };
savedObjectsMock.get.mockReset().mockResolvedValue({
id: DATA_SOURCE_ID,
type: DATA_SOURCE_SAVED_OBJECT_TYPE,
attributes: {
...mockDataSourceAttr,
auth: {
type: AuthType.SigV4,
credentials: sigV4AuthContent,
},
},
references: [],
});

jest.spyOn(cryptographyMock, 'decodeAndDecrypt').mockResolvedValue({
decryptedText: 'accessKey',
encryptionContext: { endpoint: 'http://test.com' },
});
await configureClient(dataSourceClientParams, opensearchClientPoolSetup, config, logger);

expect(ClientMock).toHaveBeenCalledTimes(2);
});
});

describe('Auth Method from Registry', () => {
beforeEach(() => {
const authMethodWithClientPool: AuthenticationMethod = {
name: 'clientPoolTest',
authType: AuthType.SigV4,
credentialProvider: jest.fn(),
clientPoolSetup: opensearchClientPoolSetup,
legacyClientPoolSetup: clientPoolSetup,
};
authenticationMethodRegistery.getAuthenticationMethod
.mockReset()
.mockImplementation(() => authMethodWithClientPool);
const mockDataSourceAttr = { ...dataSourceAttr, name: 'custom_auth' };
savedObjectsMock.get.mockReset().mockResolvedValue({
id: DATA_SOURCE_ID,
type: DATA_SOURCE_SAVED_OBJECT_TYPE,
attributes: {
...mockDataSourceAttr,
auth: {
type: AuthType.SigV4,
credentials: customAuthContent,
},
},
references: [],
});
authRegistryCredentialProviderMock.mockReturnValue({
credential: sigV4AuthContent,
type: AuthType.SigV4,
});
});
test('Auth Method from Registry: If endpoint is same for multiple requests client pool size should be 1', async () => {
await configureClient(
{ ...dataSourceClientParams, authRegistry: authenticationMethodRegistery },
clientPoolSetup,
config,
logger
);

await configureClient(
{ ...dataSourceClientParams, authRegistry: authenticationMethodRegistery },
clientPoolSetup,
config,
logger
);

expect(ClientMock).toHaveBeenCalledTimes(1);
});

test('Auth Method from Registry: If endpoint is different for two requests client pool size should be 2', async () => {
await configureClient(
{ ...dataSourceClientParams, authRegistry: authenticationMethodRegistery },
clientPoolSetup,
config,
logger
);

const mockDataSourceAttr = {
...dataSourceAttr,
endpoint: 'http://test.com',
name: 'custom_auth',
};
savedObjectsMock.get.mockReset().mockResolvedValue({
id: DATA_SOURCE_ID,
type: DATA_SOURCE_SAVED_OBJECT_TYPE,
attributes: {
...mockDataSourceAttr,
auth: {
type: AuthType.SigV4,
credentials: customAuthContent,
},
},
references: [],
});

await configureClient(
{ ...dataSourceClientParams, authRegistry: authenticationMethodRegistery },
clientPoolSetup,
config,
logger
);

expect(ClientMock).toHaveBeenCalledTimes(2);
});
});
});
});
Loading

0 comments on commit 9933938

Please sign in to comment.