From ff06ceec6c60ae3508f46442f1dda0dca10f340c Mon Sep 17 00:00:00 2001
From: NikaHsn <nika.hasani@gmail.com>
Date: Wed, 18 Sep 2024 10:28:25 -0700
Subject: [PATCH] chore(auth): update auth plugin impl and config state machine
 to use AuthOutputs instead of Auth config (#5456)

---
 .../lib/src/auth_plugin_impl.dart             |  43 +++----
 .../lib/src/model/auth_configuration.dart     | 117 ------------------
 .../machines/configuration_state_machine.dart |  47 +++----
 .../lib/common/mock_config.dart               |   4 -
 .../test/plugin/delete_user_test.dart         |   9 +-
 .../plugin/fetch_current_device_test.dart     |   5 +-
 .../test/plugin/reset_password_test.dart      |   9 +-
 ...fetch_auth_session_state_machine_test.dart |   3 +-
 8 files changed, 48 insertions(+), 189 deletions(-)
 delete mode 100644 packages/auth/amplify_auth_cognito_dart/lib/src/model/auth_configuration.dart

diff --git a/packages/auth/amplify_auth_cognito_dart/lib/src/auth_plugin_impl.dart b/packages/auth/amplify_auth_cognito_dart/lib/src/auth_plugin_impl.dart
index 651cb92dcf..f458548642 100644
--- a/packages/auth/amplify_auth_cognito_dart/lib/src/auth_plugin_impl.dart
+++ b/packages/auth/amplify_auth_cognito_dart/lib/src/auth_plugin_impl.dart
@@ -46,6 +46,8 @@ import 'package:amplify_auth_cognito_dart/src/state/state.dart';
 import 'package:amplify_auth_cognito_dart/src/util/cognito_iam_auth_provider.dart';
 import 'package:amplify_auth_cognito_dart/src/util/cognito_user_pools_auth_provider.dart';
 import 'package:amplify_core/amplify_core.dart';
+// ignore: implementation_imports
+import 'package:amplify_core/src/config/amplify_outputs/auth/auth_outputs.dart';
 // ignore: implementation_imports, invalid_use_of_internal_member
 import 'package:amplify_core/src/http/amplify_category_method.dart';
 import 'package:amplify_secure_storage_dart/amplify_secure_storage_dart.dart';
@@ -123,19 +125,15 @@ class AmplifyAuthCognitoDart extends AuthPluginInterface
     return cognitoIdp;
   }
 
-  /// The Cognito user pool configuration.
-  CognitoUserPoolConfig get _userPoolConfig {
-    final userPoolConfig = _stateMachine.get<CognitoUserPoolConfig>();
-    if (userPoolConfig == null) {
+  AuthOutputs get _authOutputs {
+    final authOutputs = _stateMachine.get<AuthOutputs>();
+    if (authOutputs?.userPoolId == null ||
+        authOutputs?.userPoolClientId == null) {
       throw const InvalidAccountTypeException.noUserPool();
     }
-    return userPoolConfig;
+    return authOutputs!;
   }
 
-  /// The Cognito identity pool configuration.
-  CognitoIdentityCredentialsProvider? get _identityPoolConfig =>
-      _stateMachine.get();
-
   /// The device metadata repository, used for handling device operations.
   DeviceMetadataRepository get _deviceRepo => _stateMachine.getOrCreate();
 
@@ -304,13 +302,13 @@ class AmplifyAuthCognitoDart extends AuthPluginInterface
   /// If there is no federation active, this is a no-op.
   /// {@endtemplate}
   Future<void> clearFederationToIdentityPool() async {
-    final identityPoolConfig = _identityPoolConfig;
-    if (identityPoolConfig == null) {
+    final identityPoolId = _authOutputs.identityPoolId;
+    if (identityPoolId == null) {
       throw const InvalidAccountTypeException.noIdentityPool();
     }
     await stateMachine.acceptAndComplete(
       CredentialStoreEvent.clearCredentials(
-        CognitoIdentityPoolKeys(identityPoolConfig.poolId),
+        CognitoIdentityPoolKeys(identityPoolId),
       ),
     );
   }
@@ -451,15 +449,16 @@ class AmplifyAuthCognitoDart extends AuthPluginInterface
     final result = await _cognitoIdp.resendConfirmationCode(
       cognito.ResendConfirmationCodeRequest.build((b) {
         b
-          ..clientId = _userPoolConfig.appClientId
+          ..clientId = _authOutputs.userPoolClientId
           ..username = username
           ..analyticsMetadata = _analyticsMetadata?.toBuilder();
 
-        final clientSecret = _userPoolConfig.appClientSecret;
+        // ignore: invalid_use_of_internal_member
+        final clientSecret = _authOutputs.appClientSecret;
         if (clientSecret != null) {
           b.secretHash = computeSecretHash(
             username,
-            _userPoolConfig.appClientId,
+            _authOutputs.userPoolClientId!,
             clientSecret,
           );
         }
@@ -763,16 +762,17 @@ class AmplifyAuthCognitoDart extends AuthPluginInterface
     final result = await _cognitoIdp.forgotPassword(
       cognito.ForgotPasswordRequest.build((b) {
         b
-          ..clientId = _userPoolConfig.appClientId
+          ..clientId = _authOutputs.userPoolClientId
           ..username = username
           ..analyticsMetadata = _analyticsMetadata?.toBuilder()
           ..clientMetadata.addAll(pluginOptions.clientMetadata);
 
-        final clientSecret = _userPoolConfig.appClientSecret;
+        // ignore: invalid_use_of_internal_member
+        final clientSecret = _authOutputs.appClientSecret;
         if (clientSecret != null) {
           b.secretHash = computeSecretHash(
             username,
-            _userPoolConfig.appClientId,
+            _authOutputs.userPoolClientId!,
             clientSecret,
           );
         }
@@ -817,15 +817,16 @@ class AmplifyAuthCognitoDart extends AuthPluginInterface
           ..username = username
           ..password = newPassword
           ..confirmationCode = confirmationCode
-          ..clientId = _userPoolConfig.appClientId
+          ..clientId = _authOutputs.userPoolClientId
           ..clientMetadata.addAll(pluginOptions.clientMetadata)
           ..analyticsMetadata = _analyticsMetadata?.toBuilder();
 
-        final clientSecret = _userPoolConfig.appClientSecret;
+        // ignore: invalid_use_of_internal_member
+        final clientSecret = _authOutputs.appClientSecret;
         if (clientSecret != null) {
           b.secretHash = computeSecretHash(
             username,
-            _userPoolConfig.appClientId,
+            _authOutputs.userPoolClientId!,
             clientSecret,
           );
         }
diff --git a/packages/auth/amplify_auth_cognito_dart/lib/src/model/auth_configuration.dart b/packages/auth/amplify_auth_cognito_dart/lib/src/model/auth_configuration.dart
deleted file mode 100644
index ca3d3f52bb..0000000000
--- a/packages/auth/amplify_auth_cognito_dart/lib/src/model/auth_configuration.dart
+++ /dev/null
@@ -1,117 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-
-@internal
-library amplify_auth_cognito.model.auth_configuration;
-
-import 'package:amplify_core/amplify_core.dart';
-import 'package:meta/meta.dart';
-
-/// {@template amplify_auth_cognito.auth_configuration}
-/// Union class for auth configurations of the different types.
-/// {@endtemplate}
-
-// TODO(nikahsn): refactor AuthConfiguration to not use AmplifyConfig types
-class AuthConfiguration with AWSEquatable<AuthConfiguration> {
-  const AuthConfiguration._({
-    this.userPoolConfig,
-    this.identityPoolConfig,
-    this.hostedUiConfig,
-    this.pinpointConfig,
-  });
-
-  /// Parses [amplifyOutputs] into the appropriate [AuthConfiguration] type.
-  ///
-  /// {@macro amplify_auth_cognito.auth_configuration}
-  factory AuthConfiguration.fromAmplifyOutputs(AmplifyOutputs amplifyOutputs) {
-    final authOutputs = amplifyOutputs.auth;
-    if (authOutputs == null) {
-      throw ConfigurationError(
-        'Invalid config: No Auth config found',
-      );
-    }
-    CognitoUserPoolConfig? userPoolConfig;
-    if (authOutputs.userPoolId != null &&
-        authOutputs.userPoolClientId != null) {
-      // ignore: invalid_use_of_internal_member
-      userPoolConfig = CognitoUserPoolConfig.fromAuthOutputs(authOutputs);
-    }
-    CognitoIdentityCredentialsProvider? identityPoolConfig;
-    if (authOutputs.identityPoolId != null) {
-      identityPoolConfig = CognitoIdentityCredentialsProvider(
-        poolId: authOutputs.identityPoolId!,
-        region: authOutputs.awsRegion,
-      );
-    }
-    CognitoOAuthConfig? hostedUiConfig;
-    if (authOutputs.userPoolClientId != null && authOutputs.oauth != null) {
-      // ignore: invalid_use_of_internal_member
-      hostedUiConfig = CognitoOAuthConfig.fromAuthOutputs(authOutputs);
-    }
-    if (userPoolConfig == null &&
-        identityPoolConfig == null &&
-        hostedUiConfig == null) {
-      throw ConfigurationError(
-        'Invalid config: No user pool or identity pool found',
-      );
-    }
-    CognitoPinpointAnalyticsConfig? pinpointConfig;
-    if (amplifyOutputs.analytics?.amazonPinpoint != null) {
-      pinpointConfig = CognitoPinpointAnalyticsConfig(
-        appId: amplifyOutputs.analytics!.amazonPinpoint!.appId,
-        region: amplifyOutputs.analytics!.amazonPinpoint!.awsRegion,
-      );
-    }
-    return AuthConfiguration._(
-      userPoolConfig: userPoolConfig,
-      identityPoolConfig: identityPoolConfig,
-      hostedUiConfig: hostedUiConfig,
-      pinpointConfig: pinpointConfig,
-    );
-  }
-
-  /// Parses [config] into the appropriate [AuthConfiguration] type.
-  ///
-  /// {@macro amplify_auth_cognito.auth_configuration}
-  factory AuthConfiguration.fromConfig(CognitoPluginConfig config) {
-    final userPoolConfig = config.cognitoUserPool?.default$;
-    final identityPoolConfig = config.credentialsProvider?.default$;
-    final hostedUiConfig = config.auth?.default$?.oAuth;
-    final pinpointConfig = config.pinpointAnalytics?.default$;
-
-    if (userPoolConfig == null &&
-        identityPoolConfig == null &&
-        hostedUiConfig == null) {
-      throw ConfigurationError(
-        'Invalid config: No user pool or identity pool found',
-      );
-    }
-
-    return AuthConfiguration._(
-      userPoolConfig: userPoolConfig,
-      identityPoolConfig: identityPoolConfig,
-      hostedUiConfig: hostedUiConfig,
-      pinpointConfig: pinpointConfig,
-    );
-  }
-
-  /// The user pool config, if available.
-  final CognitoUserPoolConfig? userPoolConfig;
-
-  /// The identity pool config, if available.
-  final CognitoIdentityCredentialsProvider? identityPoolConfig;
-
-  /// The Hosted UI config, if available.
-  final CognitoOAuthConfig? hostedUiConfig;
-
-  /// The Pinpoint config, if available.
-  final CognitoPinpointAnalyticsConfig? pinpointConfig;
-
-  @override
-  List<Object?> get props => [
-        userPoolConfig,
-        identityPoolConfig,
-        hostedUiConfig,
-        pinpointConfig,
-      ];
-}
diff --git a/packages/auth/amplify_auth_cognito_dart/lib/src/state/machines/configuration_state_machine.dart b/packages/auth/amplify_auth_cognito_dart/lib/src/state/machines/configuration_state_machine.dart
index af8eaabb55..6e0aae5185 100644
--- a/packages/auth/amplify_auth_cognito_dart/lib/src/state/machines/configuration_state_machine.dart
+++ b/packages/auth/amplify_auth_cognito_dart/lib/src/state/machines/configuration_state_machine.dart
@@ -6,13 +6,14 @@ import 'dart:async';
 // ignore: implementation_imports
 import 'package:amplify_analytics_pinpoint_dart/src/impl/analytics_client/endpoint_client/endpoint_info_store_manager.dart';
 import 'package:amplify_auth_cognito_dart/src/credentials/auth_plugin_credentials_provider.dart';
-import 'package:amplify_auth_cognito_dart/src/model/auth_configuration.dart';
 import 'package:amplify_auth_cognito_dart/src/sdk/cognito_identity.dart';
 import 'package:amplify_auth_cognito_dart/src/sdk/cognito_identity_provider.dart';
 import 'package:amplify_auth_cognito_dart/src/sdk/sdk_bridge.dart';
 import 'package:amplify_auth_cognito_dart/src/state/cognito_state_machine.dart';
 import 'package:amplify_auth_cognito_dart/src/state/state.dart';
 import 'package:amplify_core/amplify_core.dart';
+// ignore: implementation_imports
+import 'package:amplify_core/src/config/amplify_outputs/analytics/analytics_outputs.dart';
 
 /// {@template amplify_auth_cognito.configuration_state_machine}
 /// Manages configuration of the Auth category.
@@ -69,36 +70,24 @@ final class ConfigurationStateMachine
       throw ConfigurationError('No Cognito plugin config available');
     }
     addInstance(authOutputs);
-    final config = AuthConfiguration.fromAmplifyOutputs(event.config);
-    addInstance(config);
-
     final waiters = <Future<void>>[];
-    final userPoolConfig = config.userPoolConfig;
-    if (userPoolConfig != null) {
-      addInstance(userPoolConfig);
-      addInstance<CognitoIdentityProviderClient>(
-        WrappedCognitoIdentityProviderClient(
-          region: userPoolConfig.region,
-          credentialsProvider: _credentialsProvider,
-          dependencyManager: this,
-          endpoint: userPoolConfig.endpoint,
-        ),
-      );
-    }
+    addInstance<CognitoIdentityProviderClient>(
+      WrappedCognitoIdentityProviderClient(
+        region: authOutputs.awsRegion,
+        credentialsProvider: _credentialsProvider,
+        dependencyManager: this,
+      ),
+    );
 
     // Configure HostedUI, if available
-    final hostedUiConfig = config.hostedUiConfig;
-    if (hostedUiConfig != null) {
-      addInstance(hostedUiConfig);
+    if (authOutputs.oauth != null) {
       waiters.add(manager.configureHostedUI());
     }
 
-    final identityPoolConfig = config.identityPoolConfig;
-    if (identityPoolConfig != null) {
-      addInstance(identityPoolConfig);
+    if (authOutputs.identityPoolId != null) {
       addInstance<CognitoIdentityClient>(
         WrappedCognitoIdentityClient(
-          region: identityPoolConfig.region,
+          region: authOutputs.awsRegion,
           credentialsProvider: _credentialsProvider,
           dependencyManager: this,
         ),
@@ -110,7 +99,7 @@ final class ConfigurationStateMachine
     await _waitForConfiguration(event.config, waiters);
 
     // Setup AnalyticsMetadataType
-    await _registerAnalyticsMetadata(config);
+    await _registerAnalyticsMetadata(event.config.analytics);
   }
 
   Future<void> _waitForConfiguration(
@@ -124,13 +113,9 @@ final class ConfigurationStateMachine
   /// State machine callback for the [ConfigureSucceeded] event.
   Future<void> onConfigureSucceeded(ConfigureSucceeded event) async {}
 
-  Future<void> _registerAnalyticsMetadata(AuthConfiguration config) async {
-    final analyticsConfig = config.pinpointConfig;
-    if (analyticsConfig == null) {
-      return;
-    }
-    final appId = analyticsConfig.appId;
-
+  Future<void> _registerAnalyticsMetadata(AnalyticsOutputs? config) async {
+    final appId = config?.amazonPinpoint?.appId;
+    if (appId == null) return;
     final endpointStoreManager = getOrCreate<EndpointInfoStoreManager>();
     await endpointStoreManager.init(pinpointAppId: appId);
 
diff --git a/packages/auth/amplify_auth_cognito_test/lib/common/mock_config.dart b/packages/auth/amplify_auth_cognito_test/lib/common/mock_config.dart
index 1288d86930..b97d31348a 100644
--- a/packages/auth/amplify_auth_cognito_test/lib/common/mock_config.dart
+++ b/packages/auth/amplify_auth_cognito_test/lib/common/mock_config.dart
@@ -6,7 +6,6 @@ import 'dart:convert';
 import 'package:amplify_auth_cognito_dart/amplify_auth_cognito_dart.dart';
 import 'package:amplify_auth_cognito_dart/src/credentials/cognito_keys.dart';
 import 'package:amplify_auth_cognito_dart/src/jwt/jwt.dart';
-import 'package:amplify_auth_cognito_dart/src/model/auth_configuration.dart';
 import 'package:amplify_core/amplify_core.dart';
 
 const amplifyConfig = '''{
@@ -131,9 +130,6 @@ final mockConfigWithPinpoint = AmplifyOutputs.fromJson(
   jsonDecode(amplifyConfigWithAnalytics) as Map<String, Object?>,
 );
 
-final authConfig = AuthConfiguration.fromAmplifyOutputs(mockConfig);
-// TODO(nikahsn): remove after refactoring DeviceMetadataRepository
-final userPoolConfig = authConfig.userPoolConfig!;
 final userPoolKeys = CognitoUserPoolKeys(mockConfig.auth!.userPoolClientId!);
 final deviceKeys =
     CognitoDeviceKeys(mockConfig.auth!.userPoolClientId!, userSub);
diff --git a/packages/auth/amplify_auth_cognito_test/test/plugin/delete_user_test.dart b/packages/auth/amplify_auth_cognito_test/test/plugin/delete_user_test.dart
index ded6e2d45f..b7944eaabb 100644
--- a/packages/auth/amplify_auth_cognito_test/test/plugin/delete_user_test.dart
+++ b/packages/auth/amplify_auth_cognito_test/test/plugin/delete_user_test.dart
@@ -8,7 +8,6 @@ import 'dart:async';
 import 'package:amplify_auth_cognito_dart/amplify_auth_cognito_dart.dart'
     hide InternalErrorException;
 import 'package:amplify_auth_cognito_dart/src/credentials/cognito_keys.dart';
-import 'package:amplify_auth_cognito_dart/src/model/auth_configuration.dart';
 import 'package:amplify_auth_cognito_dart/src/sdk/cognito_identity_provider.dart';
 import 'package:amplify_auth_cognito_dart/src/state/cognito_state_machine.dart';
 import 'package:amplify_auth_cognito_test/common/matchers.dart';
@@ -20,11 +19,9 @@ import 'package:amplify_secure_storage_dart/amplify_secure_storage_dart.dart';
 import 'package:test/test.dart';
 
 void main() {
-  final authConfig = AuthConfiguration.fromAmplifyOutputs(mockConfig);
-  final userPoolConfig = authConfig.userPoolConfig!;
-  final identityPoolConfig = authConfig.identityPoolConfig!;
-  final userPoolKeys = CognitoUserPoolKeys(userPoolConfig.appClientId);
-  final identityPoolKeys = CognitoIdentityPoolKeys(identityPoolConfig.poolId);
+  final userPoolKeys = CognitoUserPoolKeys(mockConfig.auth!.userPoolClientId!);
+  final identityPoolKeys =
+      CognitoIdentityPoolKeys(mockConfig.auth!.identityPoolId!);
 
   late AmplifyAuthCognitoDart plugin;
   late CognitoAuthStateMachine stateMachine;
diff --git a/packages/auth/amplify_auth_cognito_test/test/plugin/fetch_current_device_test.dart b/packages/auth/amplify_auth_cognito_test/test/plugin/fetch_current_device_test.dart
index d2852b8a13..e59daf6985 100644
--- a/packages/auth/amplify_auth_cognito_test/test/plugin/fetch_current_device_test.dart
+++ b/packages/auth/amplify_auth_cognito_test/test/plugin/fetch_current_device_test.dart
@@ -14,7 +14,7 @@ import 'package:test/test.dart';
 void main() {
   AmplifyLogger().logLevel = LogLevel.verbose;
 
-  final userPoolKeys = CognitoUserPoolKeys(userPoolConfig.appClientId);
+  final userPoolKeys = CognitoUserPoolKeys(mockConfig.auth!.userPoolClientId!);
   final identityPoolKeys =
       CognitoIdentityPoolKeys(mockConfig.auth!.identityPoolId!);
   final testAuthRepo = AmplifyAuthProviderRepository();
@@ -31,7 +31,8 @@ void main() {
         secureStorage,
         userPoolKeys: userPoolKeys,
         identityPoolKeys: identityPoolKeys,
-        deviceKeys: CognitoDeviceKeys(userPoolConfig.appClientId, username),
+        deviceKeys:
+            CognitoDeviceKeys(mockConfig.auth!.userPoolClientId!, username),
       );
       plugin = AmplifyAuthCognitoDart(
         secureStorageFactory: (_) => secureStorage,
diff --git a/packages/auth/amplify_auth_cognito_test/test/plugin/reset_password_test.dart b/packages/auth/amplify_auth_cognito_test/test/plugin/reset_password_test.dart
index 69a305b113..60beb45545 100644
--- a/packages/auth/amplify_auth_cognito_test/test/plugin/reset_password_test.dart
+++ b/packages/auth/amplify_auth_cognito_test/test/plugin/reset_password_test.dart
@@ -3,7 +3,6 @@
 
 import 'package:amplify_auth_cognito_dart/amplify_auth_cognito_dart.dart';
 import 'package:amplify_auth_cognito_dart/src/credentials/cognito_keys.dart';
-import 'package:amplify_auth_cognito_dart/src/model/auth_configuration.dart';
 import 'package:amplify_auth_cognito_dart/src/sdk/cognito_identity_provider.dart'
     hide NotAuthorizedException;
 import 'package:amplify_auth_cognito_dart/src/state/cognito_state_machine.dart';
@@ -17,11 +16,9 @@ import 'package:test/test.dart';
 // Follows resetPassword test cases:
 // https://github.com/aws-amplify/amplify-android/tree/main/aws-auth-cognito/src/test/resources/feature-test/testsuites/resetPassword
 void main() {
-  final authConfig = AuthConfiguration.fromAmplifyOutputs(mockConfig);
-  final userPoolConfig = authConfig.userPoolConfig!;
-  final identityPoolConfig = authConfig.identityPoolConfig!;
-  final userPoolKeys = CognitoUserPoolKeys(userPoolConfig.appClientId);
-  final identityPoolKeys = CognitoIdentityPoolKeys(identityPoolConfig.poolId);
+  final userPoolKeys = CognitoUserPoolKeys(mockConfig.auth!.userPoolClientId!);
+  final identityPoolKeys =
+      CognitoIdentityPoolKeys(mockConfig.auth!.identityPoolId!);
 
   late AmplifyAuthCognitoDart plugin;
   late CognitoAuthStateMachine stateMachine;
diff --git a/packages/auth/amplify_auth_cognito_test/test/state/fetch_auth_session_state_machine_test.dart b/packages/auth/amplify_auth_cognito_test/test/state/fetch_auth_session_state_machine_test.dart
index ce7f2747bc..fe819884b8 100644
--- a/packages/auth/amplify_auth_cognito_test/test/state/fetch_auth_session_state_machine_test.dart
+++ b/packages/auth/amplify_auth_cognito_test/test/state/fetch_auth_session_state_machine_test.dart
@@ -120,8 +120,7 @@ void main() {
       secureStorage = MockSecureStorage();
       stateMachine = CognitoAuthStateMachine()
         ..addInstance(secureStorage)
-        ..addInstance(mockConfig)
-        ..addInstance(authConfig);
+        ..addInstance(mockConfig);
     });
 
     group('User Pool + Identity Pool', () {