From 14db3b4e25ca551396dbf3a351381e1e1a426600 Mon Sep 17 00:00:00 2001
From: Dillon Nys <24740863+dnys1@users.noreply.github.com>
Date: Wed, 6 Sep 2023 08:32:37 -0700
Subject: [PATCH] test(auth): Fix flaky test (#3678)

Fixes flaky HostedUI test by updating expectations and not relying on the timings of the event loop.
---
 .../hosted_ui/hosted_ui_platform_io.dart      | 29 ++++------
 .../hostedui/hosted_ui_platform_io_test.dart  | 53 ++++++++++---------
 2 files changed, 39 insertions(+), 43 deletions(-)

diff --git a/packages/auth/amplify_auth_cognito_dart/lib/src/flows/hosted_ui/hosted_ui_platform_io.dart b/packages/auth/amplify_auth_cognito_dart/lib/src/flows/hosted_ui/hosted_ui_platform_io.dart
index 28791f6aca..aca34f8f5c 100644
--- a/packages/auth/amplify_auth_cognito_dart/lib/src/flows/hosted_ui/hosted_ui_platform_io.dart
+++ b/packages/auth/amplify_auth_cognito_dart/lib/src/flows/hosted_ui/hosted_ui_platform_io.dart
@@ -156,8 +156,7 @@ class HostedUiPlatformImpl extends HostedUiPlatform {
   /// of the ports may be blocked.
   @visibleForTesting
   Future<LocalServer> localConnect(Iterable<Uri> uris) async {
-    final localServer = _localServer;
-    if (localServer != null) {
+    if (_localServer case final localServer?) {
       return localServer;
     }
 
@@ -182,7 +181,7 @@ class HostedUiPlatformImpl extends HostedUiPlatform {
             'All ports were blocked: [${uris.map((uri) => uri.port).join(',')}]',
       );
     }
-    return LocalServer(server, selectedUri);
+    return _localServer = LocalServer(server, selectedUri);
   }
 
   Future<void> _respond(
@@ -212,20 +211,16 @@ class HostedUiPlatformImpl extends HostedUiPlatform {
       _noSuitableRedirect(signIn: true);
     }
 
-    _localServer = await localConnect(signInUris);
+    final localServer = await localConnect(signInUris);
     try {
       final signInUrl = (await getSignInUri(
         provider: provider,
-        redirectUri: _localServer!.uri,
+        redirectUri: localServer.uri,
       ))
           .toString();
       await launchUrl(signInUrl);
 
-      final server = _localServer?.server;
-      if (server == null) {
-        return;
-      }
-      await for (final request in server) {
+      await for (final request in localServer.server) {
         final method = request.method;
         if (method != 'GET') {
           await _respond(
@@ -288,17 +283,12 @@ class HostedUiPlatformImpl extends HostedUiPlatform {
       _noSuitableRedirect(signIn: false);
     }
 
-    _localServer = await localConnect(signOutUris);
+    final localServer = await localConnect(signOutUris);
     try {
-      final signOutUri =
-          getSignOutUri(redirectUri: _localServer!.uri).toString();
+      final signOutUri = getSignOutUri(redirectUri: localServer.uri).toString();
       await launchUrl(signOutUri);
 
-      final server = _localServer?.server;
-      if (server == null) {
-        return;
-      }
-      await for (final request in server) {
+      await for (final request in localServer.server) {
         final method = request.method;
         if (method != 'GET') {
           await _respond(
@@ -331,8 +321,9 @@ class HostedUiPlatformImpl extends HostedUiPlatform {
   /// Closes the open server, if any.
   @override
   Future<void> close() async {
-    await _localServer?.server.close();
+    final localServer = _localServer;
     _localServer = null;
+    await localServer?.server.close();
   }
 }
 
diff --git a/packages/auth/amplify_auth_cognito_test/test/flows/hostedui/hosted_ui_platform_io_test.dart b/packages/auth/amplify_auth_cognito_test/test/flows/hostedui/hosted_ui_platform_io_test.dart
index 583b53571e..6673d1cf6a 100644
--- a/packages/auth/amplify_auth_cognito_test/test/flows/hostedui/hosted_ui_platform_io_test.dart
+++ b/packages/auth/amplify_auth_cognito_test/test/flows/hostedui/hosted_ui_platform_io_test.dart
@@ -21,20 +21,16 @@ import 'package:test/test.dart';
 class MockHostedUiPlatform extends HostedUiPlatformImpl {
   MockHostedUiPlatform(super.dependencyManager);
 
-  LocalServer? _localServer;
+  final launchedServer = Completer<void>();
 
   @override
   Future<void> launchUrl(String url) async {}
 
   @override
   Future<LocalServer> localConnect(Iterable<Uri> uris) async {
-    return _localServer = await super.localConnect(uris);
-  }
-
-  @override
-  Future<void> close() async {
-    _localServer = null;
-    await super.close();
+    final localServer = await super.localConnect(uris);
+    launchedServer.complete();
+    return localServer;
   }
 }
 
@@ -153,32 +149,41 @@ void main() {
         expect(hostedUiPlatform.signInRedirectUri, redirect);
         expect(hostedUiPlatform.signOutRedirectUri, redirect);
 
-        unawaited(
+        expect(
           hostedUiPlatform.signIn(
             options: const CognitoSignInWithWebUIPluginOptions(),
           ),
+          completes,
         );
 
-        await expectLater(client.get(redirect), completes);
-        expect(
-          hostedUiPlatform._localServer,
-          isNotNull,
-          reason: "Server won't close until a valid redirect is performed",
+        await hostedUiPlatform.launchedServer.future;
+
+        await expectLater(
+          client.get(redirect),
+          completion(
+            isA<http.Response>()
+                .having((res) => res.statusCode, 'statusCode', isNot(200)),
+          ),
+          reason: 'Local server should be running and able to accept '
+              'requests. Should return 4xx for anything but a success/error '
+              'redirect.',
         );
 
+        final successRedirect = redirect.replace(
+          queryParameters: {'state': 'state', 'code': 'code'},
+        );
         await expectLater(
-          client.get(
-            redirect.replace(
-              queryParameters: {'state': 'state', 'code': 'code'},
-            ),
+          client.get(successRedirect),
+          completion(
+            isA<http.Response>()
+                .having((res) => res.statusCode, 'statusCode', 200),
           ),
-          completes,
+          reason: "Server won't close until a valid redirect is performed",
         );
-        await Future<void>.delayed(Duration.zero);
-        expect(
-          hostedUiPlatform._localServer,
-          isNull,
-          reason: 'A valid redirect includes state + code/error',
+        await expectLater(
+          client.get(redirect),
+          throwsA(isA<http.ClientException>()),
+          reason: 'Server should be closed after successful redirect',
         );
       });
     });