feat(ui_auth): add a way to get an AuthController for AuthState

packages/firebase_ui_auth/lib/src/actions.dart

@@ -180,7 +180,9 @@ class _FlutterfireUIAuthActionsElement extends InheritedElement {
       listener: (oldState, newState, controller) {
         for (final action in widget.actions) {
           if (action is AuthStateChangeAction && action.matches(newState)) {
+            _controllerRegistry[newState] = controller;
             action.invoke(this, newState);
+            _controllerRegistry.remove(newState);
           }
         }
 
@@ -190,3 +192,38 @@ class _FlutterfireUIAuthActionsElement extends InheritedElement {
     );
   }
 }
+
+final _controllerRegistry = <AuthState, AuthController>{};
+
+/// Allows getting an [AuthController] that caused an [AuthState] transition.
+/// Calling [getControllerForState] is only allowed from a [FirebaseUIAction]
+/// callback:
+///
+/// ```dart
+/// SignInScreen(
+///   actions: [
+///     AuthStateChangeAction<SignedIn>((context, state) {
+///       final ctrl = getControllerForState(state);
+///
+///       if (ctrl is EmailAuthController) {
+///          print('email was used for authentication');
+///       }
+///
+///
+///       Navigator.of(context).pushReplacementNamed('/profile');
+///     }
+///   ]
+/// );
+/// ```
+AuthController getControllerForState(AuthState state) {
+  final ctrl = _controllerRegistry[state];
+
+  if (ctrl == null) {
+    throw StateError(
+      'Quering controller for an auth state is only allowed '
+      'from FirebaseUIAction callback',
+    );
+  }
+
+  return ctrl;
+}

tests/integration_test/firebase_ui_auth/actions_test.dart

@@ -0,0 +1,78 @@
+// Copyright 2023, the Chromium project authors.  Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'package:firebase_ui_auth/firebase_ui_auth.dart';
+import 'package:flutter/material.dart';
+import 'package:flutter_test/flutter_test.dart';
+
+import '../utils.dart';
+
+void main() {
+  setUpTests();
+  group('getControllerForType', () {
+    setUp(() async {
+      await auth.createUserWithEmailAndPassword(
+        email: '[email protected]',
+        password: '123456',
+      );
+    });
+
+    Future<void> authenticate(WidgetTester tester) async {
+      final inputs = find.byType(TextFormField);
+
+      await tester.enterText(inputs.at(0), '[email protected]');
+      await tester.testTextInput.receiveAction(TextInputAction.done);
+      await tester.pumpAndSettle();
+
+      await tester.enterText(inputs.at(1), '123456');
+      await tester.testTextInput.receiveAction(TextInputAction.done);
+      await tester.pumpAndSettle();
+    }
+
+    testWidgets(
+      'getControllerForState returns correct controller type',
+      (tester) async {
+        late AuthController controller;
+
+        await render(
+          tester,
+          FirebaseUIActions(
+            actions: [
+              AuthStateChangeAction<SignedIn>((context, state) {
+                controller = getControllerForState(state);
+              })
+            ],
+            child: const EmailForm(action: AuthAction.signIn),
+          ),
+        );
+
+        await authenticate(tester);
+
+        expect(controller, isA<EmailAuthController>());
+      },
+    );
+
+    testWidgets(
+      'throws a StateError if used outside of the FirebaseUIAction',
+      (tester) async {
+        late AuthState state;
+
+        await render(
+          tester,
+          AuthStateListener(
+            listener: (oldState, newState, _) {
+              state = newState;
+              return null;
+            },
+            child: const EmailForm(action: AuthAction.signIn),
+          ),
+        );
+
+        await authenticate(tester);
+
+        expect(() => getControllerForState(state), throwsStateError);
+      },
+    );
+  });
+}

tests/integration_test/firebase_ui_auth/firebase_ui_auth_e2e.dart

@@ -11,12 +11,14 @@ import 'universal_email_sign_in_screen_test.dart'
     as universal_email_sign_in_screen;
 import 'phone_verification_test.dart' as phone_verification;
 import 'layout_test.dart' as layout;
+import 'actions_test.dart' as actions;
 
 Future<void> main() async {
   group('Auth', () {
     email_form.main();
     email_link_sign_in_view.main();
     universal_email_sign_in_screen.main();
+    actions.main();
 
     switch (defaultTargetPlatform) {
       case TargetPlatform.iOS:

tests/integration_test/firebase_ui_test.dart

@@ -6,6 +6,7 @@ import 'package:flutter/foundation.dart';
 import 'package:flutter_test/flutter_test.dart';
 
 import './firebase_ui_auth/firebase_ui_auth_e2e.dart' as firebase_ui_auth_e2e;
+
 import './firebase_ui_firestore/firebase_ui_firestore_e2e.dart'
     as firebase_ui_firestore_e2e;
 import './firebase_ui_oauth_apple/firebase_ui_oauth_apple_e2e.dart'
@@ -20,10 +21,8 @@ import './firebase_ui_oauth_twitter/firebase_ui_oauth_twitter_e2e.dart'
 import 'utils.dart';
 
 void main() {
+  setUpTests();
   group('Firebase UI', () {
-    setUpAll(prepare);
-    tearDown(authCleanup);
-
     firebase_ui_auth_e2e.main();
 
     if (defaultTargetPlatform != TargetPlatform.macOS) {

tests/integration_test/utils.dart

@@ -58,6 +58,16 @@ Future<void> authCleanup() async {
   await deleteAllAccounts();
 }
 
+bool _testsSetUp = false;
+
+void setUpTests() {
+  if (_testsSetUp) return;
+
+  setUpAll(prepare);
+  tearDown(authCleanup);
+  _testsSetUp = true;
+}
+
 Future<void> render(WidgetTester tester, Widget widget) async {
   await tester.pumpWidget(
     MaterialApp(