From 0191532a1a8dc6b7280ae883973c77ca24ab1f3c Mon Sep 17 00:00:00 2001
From: Di Wu <im@wudi.me>
Date: Fri, 3 Nov 2023 14:55:06 -0700
Subject: [PATCH] feat(core): add API for partially config Amplify

---
 .../__tests__/singleton/Singleton.test.ts     | 71 ++++++++++++++++++-
 packages/core/src/singleton/Amplify.ts        | 18 +++++
 packages/core/src/utils/Lens.ts               | 35 +++++++++
 3 files changed, 123 insertions(+), 1 deletion(-)
 create mode 100644 packages/core/src/utils/Lens.ts

diff --git a/packages/core/__tests__/singleton/Singleton.test.ts b/packages/core/__tests__/singleton/Singleton.test.ts
index a6a51f8c068..00a02e11e2f 100644
--- a/packages/core/__tests__/singleton/Singleton.test.ts
+++ b/packages/core/__tests__/singleton/Singleton.test.ts
@@ -3,7 +3,8 @@ import { AuthClass as Auth } from '../../src/singleton/Auth';
 import { decodeJWT } from '../../src/singleton/Auth/utils';
 import { AWSCredentialsAndIdentityId } from '../../src/singleton/Auth/types';
 import { TextEncoder, TextDecoder } from 'util';
-import { fetchAuthSession } from '../../src';
+import { fetchAuthSession, ResourcesConfig } from '../../src';
+import { extend } from 'lodash';
 Object.assign(global, { TextDecoder, TextEncoder });
 type ArgumentTypes<F extends Function> = F extends (...args: infer A) => any
 	? A
@@ -17,6 +18,74 @@ const MOCK_AUTH_CONFIG = {
 	},
 };
 
+describe('Partial configure', () => {
+	it('should update partial config', () => {
+		const mockConfig: ResourcesConfig = {
+			Auth: {
+				Cognito: {
+					allowGuestAccess: true,
+					identityPoolId: 'aws_cognito_identity_pool_id',
+					userPoolClientId: 'aws_user_pools_web_client_id',
+					userPoolId: 'aws_user_pools_id',
+					passwordFormat: {
+						minLength: 8,
+						requireLowercase: false,
+						requireNumbers: false,
+						requireSpecialCharacters: false,
+						requireUppercase: false,
+					},
+				},
+			},
+		};
+
+		Amplify.configure(mockConfig);
+		Amplify.updateConfig(
+			config =>
+				config
+					.atKey('Auth')
+					.atKey('Cognito')
+					.atKey('passwordFormat')
+					.atKey('minLength'),
+			12
+		);
+		expect(Amplify.getConfig()).toStrictEqual({
+			Auth: {
+				Cognito: {
+					allowGuestAccess: true,
+					identityPoolId: 'aws_cognito_identity_pool_id',
+					userPoolClientId: 'aws_user_pools_web_client_id',
+					userPoolId: 'aws_user_pools_id',
+					passwordFormat: {
+						minLength: 12,
+						requireLowercase: false,
+						requireNumbers: false,
+						requireSpecialCharacters: false,
+						requireUppercase: false,
+					},
+				},
+			},
+		});
+	});
+
+	it('should fail if update non existing field', () => {
+		const mockConfig: ResourcesConfig = {
+			Analytics: {
+				Pinpoint: {
+					appId: '123',
+					region: 'us-east-1',
+				},
+			},
+		};
+		Amplify.configure(mockConfig);
+		expect(() =>
+			Amplify.updateConfig(
+				config => config.atKey('Auth').atKey('Cognito').atKey('identityPoolId'),
+				'0000'
+			)
+		).toThrow();
+	});
+});
+
 describe('Amplify.configure() and Amplify.getConfig()', () => {
 	it('should take the legacy CLI shaped config object for configuring and return it from getConfig()', () => {
 		const mockLegacyConfig = {
diff --git a/packages/core/src/singleton/Amplify.ts b/packages/core/src/singleton/Amplify.ts
index 3040c95c8b6..606469fb13c 100644
--- a/packages/core/src/singleton/Amplify.ts
+++ b/packages/core/src/singleton/Amplify.ts
@@ -5,6 +5,7 @@ import { Hub, AMPLIFY_SYMBOL } from '../Hub';
 import { LegacyConfig, LibraryOptions, ResourcesConfig } from './types';
 import { parseAWSExports } from '../parseAWSExports';
 import { deepFreeze } from '../utils';
+import { Lens } from '../utils/Lens';
 
 export class AmplifyClass {
 	resourcesConfig: ResourcesConfig;
@@ -78,6 +79,23 @@ export class AmplifyClass {
 	getConfig(): Readonly<ResourcesConfig> {
 		return this.resourcesConfig;
 	}
+
+	/**
+	 * Partially update the configuration
+	 *
+	 * @param selector Selector to define which field to set.
+	 * @param value The value to set.
+	 * */
+	updateConfig<T>(
+		selector: (
+			config: Lens<ResourcesConfig, ResourcesConfig>
+		) => Lens<ResourcesConfig, T>,
+		value: T
+	): void {
+		Amplify.configure(
+			selector(Lens.of<ResourcesConfig>()).set(value)(this.resourcesConfig)
+		);
+	}
 }
 
 /**
diff --git a/packages/core/src/utils/Lens.ts b/packages/core/src/utils/Lens.ts
new file mode 100644
index 00000000000..0de6eef5d7f
--- /dev/null
+++ b/packages/core/src/utils/Lens.ts
@@ -0,0 +1,35 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+type Getter<W, P> = (w: W) => P;
+type Setter<W, P> = (p: P) => (w: W) => W;
+
+export class Lens<W, P> {
+	private readonly get: Getter<W, P>;
+
+	/** @internal */
+	readonly set: Setter<W, P>;
+
+	static of<T>(): Lens<T, T> {
+		return new Lens<T, T>(
+			t => t,
+			t => () => t
+		);
+	}
+
+	constructor(get: Getter<W, P>, set: Setter<W, P>) {
+		this.get = get;
+		this.set = set;
+	}
+
+	atKey<Q extends keyof P>(key: Q): Lens<W, P[Q]> {
+		return new Lens<W, P[Q]>(
+			w => this.get(w)[key],
+			pq => w =>
+				this.set({
+					...this.get(w),
+					[key]: pq,
+				})(w)
+		);
+	}
+}