Use React Suspense when rendering async modals (#28386)

* Use React Suspense when rendering async modals

Signed-off-by: Michael Telatynski <7t3chguy@gmail.com>

* Fix test

Signed-off-by: Michael Telatynski <7t3chguy@gmail.com>

* Improve coverage

Signed-off-by: Michael Telatynski <7t3chguy@gmail.com>

* Improve coverage

Signed-off-by: Michael Telatynski <7t3chguy@gmail.com>

* Improve coverage

Signed-off-by: Michael Telatynski <7t3chguy@gmail.com>

* Update src/Modal.tsx

---------

Signed-off-by: Michael Telatynski <7t3chguy@gmail.com>

src/AddThreepid.ts

@@ -10,13 +10,13 @@ Please see LICENSE files in the repository root for full details.
 
 import {
     IAddThreePidOnlyBody,
-    IAuthData,
     IRequestMsisdnTokenResponse,
     IRequestTokenResponse,
     MatrixClient,
     MatrixError,
     HTTPError,
     IThreepid,
+    UIAResponse,
 } from "matrix-js-sdk/src/matrix";
 
 import Modal from "./Modal";
@@ -179,7 +179,9 @@ export default class AddThreepid {
      * with a "message" property which contains a human-readable message detailing why
      * the request failed.
      */
-    public async checkEmailLinkClicked(): Promise<[success?: boolean, result?: IAuthData | Error | null]> {
+    public async checkEmailLinkClicked(): Promise<
+        [success?: boolean, result?: UIAResponse<IAddThreePidOnlyBody> | Error | null]
+    > {
         try {
             if (this.bind) {
                 const authClient = new IdentityAuthClient();
@@ -220,7 +222,7 @@ export default class AddThreepid {
                             continueKind: "primary",
                         },
                     };
-                    const { finished } = Modal.createDialog(InteractiveAuthDialog<{}>, {
+                    const { finished } = Modal.createDialog(InteractiveAuthDialog<IAddThreePidOnlyBody>, {
                         title: _t("settings|general|add_email_dialog_title"),
                         matrixClient: this.matrixClient,
                         authData: err.data,
@@ -263,7 +265,9 @@ export default class AddThreepid {
      * with a "message" property which contains a human-readable message detailing why
      * the request failed.
      */
-    public async haveMsisdnToken(msisdnToken: string): Promise<[success?: boolean, result?: IAuthData | Error | null]> {
+    public async haveMsisdnToken(
+        msisdnToken: string,
+    ): Promise<[success?: boolean, result?: UIAResponse<IAddThreePidOnlyBody> | Error | null]> {
         const authClient = new IdentityAuthClient();
 
         if (this.submitUrl) {
@@ -319,7 +323,7 @@ export default class AddThreepid {
                         continueKind: "primary",
                     },
                 };
-                const { finished } = Modal.createDialog(InteractiveAuthDialog<{}>, {
+                const { finished } = Modal.createDialog(InteractiveAuthDialog<IAddThreePidOnlyBody>, {
                     title: _t("settings|general|add_msisdn_dialog_title"),
                     matrixClient: this.matrixClient,
                     authData: err.data,

src/AsyncWrapper.tsx

@@ -6,24 +6,19 @@ SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only
 Please see LICENSE files in the repository root for full details.
 */
 
-import React, { ComponentType, PropsWithChildren } from "react";
+import React, { ReactNode, Suspense } from "react";
-import { logger } from "matrix-js-sdk/src/logger";
 
 import { _t } from "./languageHandler";
 import BaseDialog from "./components/views/dialogs/BaseDialog";
 import DialogButtons from "./components/views/elements/DialogButtons";
 import Spinner from "./components/views/elements/Spinner";
 
-type AsyncImport<T> = { default: T };
-
 interface IProps {
-    // A promise which resolves with the real component
-    prom: Promise<ComponentType<any> | AsyncImport<ComponentType<any>>>;
     onFinished(): void;
+    children: ReactNode;
 }
 
 interface IState {
-    component?: ComponentType<PropsWithChildren<any>>;
     error?: Error;
 }
 
@@ -32,56 +27,26 @@ interface IState {
  * spinner until the real component loads.
  */
 export default class AsyncWrapper extends React.Component<IProps, IState> {
-    private unmounted = false;
+    public static getDerivedStateFromError(error: Error): IState {
+        return { error };
+    }
 
     public state: IState = {};
 
-    public componentDidMount(): void {
-        this.unmounted = false;
-        this.props.prom
-            .then((result) => {
-                if (this.unmounted) return;
-
-                // Take the 'default' member if it's there, then we support
-                // passing in just an import()ed module, since ES6 async import
-                // always returns a module *namespace*.
-                const component = (result as AsyncImport<ComponentType>).default
-                    ? (result as AsyncImport<ComponentType>).default
-                    : (result as ComponentType);
-                this.setState({ component });
-            })
-            .catch((e) => {
-                logger.warn("AsyncWrapper promise failed", e);
-                this.setState({ error: e });
-            });
-    }
-
-    public componentWillUnmount(): void {
-        this.unmounted = true;
-    }
-
-    private onWrapperCancelClick = (): void => {
-        this.props.onFinished();
-    };
-
     public render(): React.ReactNode {
-        if (this.state.component) {
+        if (this.state.error) {
-            const Component = this.state.component;
-            return <Component {...this.props} />;
-        } else if (this.state.error) {
             return (
                 <BaseDialog onFinished={this.props.onFinished} title={_t("common|error")}>
                     {_t("failed_load_async_component")}
                     <DialogButtons
                         primaryButton={_t("action|dismiss")}
-                        onPrimaryButtonClick={this.onWrapperCancelClick}
+                        onPrimaryButtonClick={this.props.onFinished}
                         hasCancel={false}
                     />
                 </BaseDialog>
             );
-        } else {
-            // show a spinner until the component is loaded.
-            return <Spinner />;
         }
+
+        return <Suspense fallback={<Spinner />}>{this.props.children}</Suspense>;
     }
 }

src/Modal.tsx

@@ -136,32 +136,6 @@ export class ModalManager extends TypedEventEmitter<ModalManagerEvent, HandlerMa
         return !!this.priorityModal || !!this.staticModal || this.modals.length > 0;
     }
 
-    public createDialog<C extends ComponentType>(
-        Element: C,
-        props?: ComponentProps<C>,
-        className?: string,
-        isPriorityModal = false,
-        isStaticModal = false,
-        options: IOptions<C> = {},
-    ): IHandle<C> {
-        return this.createDialogAsync<C>(
-            Promise.resolve(Element),
-            props,
-            className,
-            isPriorityModal,
-            isStaticModal,
-            options,
-        );
-    }
-
-    public appendDialog<C extends ComponentType>(
-        Element: C,
-        props?: ComponentProps<C>,
-        className?: string,
-    ): IHandle<C> {
-        return this.appendDialogAsync<C>(Promise.resolve(Element), props, className);
-    }
-
     /**
      * DEPRECATED.
      * This is used only for tests. They should be using forceCloseAllModals but that
@@ -196,8 +170,11 @@ export class ModalManager extends TypedEventEmitter<ModalManagerEvent, HandlerMa
         this.reRender();
     }
 
+    /**
+     * @typeParam C - the component type
+     */
     private buildModal<C extends ComponentType>(
-        prom: Promise<C>,
+        Component: C,
         props?: ComponentProps<C>,
         className?: string,
         options?: IOptions<C>,
@@ -222,9 +199,12 @@ export class ModalManager extends TypedEventEmitter<ModalManagerEvent, HandlerMa
         // otherwise we'll get confused.
         const modalCount = this.counter++;
 
-        // FIXME: If a dialog uses getDefaultProps it clobbers the onFinished
+        // Typescript doesn't like us passing props as any here, but we know that they are well typed due to the rigorous generics.
-        // property set here so you can't close the dialog from a button click!
+        modal.elem = (
-        modal.elem = <AsyncWrapper key={modalCount} prom={prom} {...props} onFinished={closeDialog} />;
+            <AsyncWrapper key={modalCount} onFinished={closeDialog}>
+                <Component {...(props as any)} onFinished={closeDialog} />
+            </AsyncWrapper>
+        );
         modal.close = closeDialog;
 
         return { modal, closeDialog, onFinishedProm };
@@ -291,29 +271,30 @@ export class ModalManager extends TypedEventEmitter<ModalManagerEvent, HandlerMa
      *       require(['<module>'], cb);
      *   }
      *
-     * @param {Promise} prom   a promise which resolves with a React component
+     * @param component The component to render as a dialog. This component must accept an `onFinished` prop function as
-     *   which will be displayed as the modal view.
+     *                  per the type {@link ComponentType}. If loading a component with esoteric dependencies consider
+     *                  using React.lazy to async load the component.
+     *                  e.g. `lazy(() => import('./MyComponent'))`
      *
-     * @param {Object} props   properties to pass to the displayed
+     * @param props properties to pass to the displayed component. (We will also pass an 'onFinished' property.)
-     *    component. (We will also pass an 'onFinished' property.)
      *
-     * @param {String} className   CSS class to apply to the modal wrapper
+     * @param className CSS class to apply to the modal wrapper
      *
-     * @param {boolean} isPriorityModal if true, this modal will be displayed regardless
+     * @param isPriorityModal if true, this modal will be displayed regardless
      *                                  of other modals that are currently in the stack.
      *                                  Also, when closed, all modals will be removed
      *                                  from the stack.
-     * @param {boolean} isStaticModal  if true, this modal will be displayed under other
+     * @param isStaticModal if true, this modal will be displayed under other
      *                                 modals in the stack. When closed, all modals will
      *                                 also be removed from the stack. This is not compatible
      *                                 with being a priority modal. Only one modal can be
      *                                 static at a time.
-     * @param {Object} options? extra options for the dialog
+     * @param options? extra options for the dialog
-     * @param {onBeforeClose} options.onBeforeClose a callback to decide whether to close the dialog
+     * @param options.onBeforeClose a callback to decide whether to close the dialog
-     * @returns {object} Object with 'close' parameter being a function that will close the dialog
+     * @returns Object with 'close' parameter being a function that will close the dialog
      */
-    public createDialogAsync<C extends ComponentType>(
+    public createDialog<C extends ComponentType>(
-        prom: Promise<C>,
+        component: C,
         props?: ComponentProps<C>,
         className?: string,
         isPriorityModal = false,
@@ -321,7 +302,7 @@ export class ModalManager extends TypedEventEmitter<ModalManagerEvent, HandlerMa
         options: IOptions<C> = {},
     ): IHandle<C> {
         const beforeModal = this.getCurrentModal();
-        const { modal, closeDialog, onFinishedProm } = this.buildModal<C>(prom, props, className, options);
+        const { modal, closeDialog, onFinishedProm } = this.buildModal<C>(component, props, className, options);
         if (isPriorityModal) {
             // XXX: This is destructive
             this.priorityModal = modal;
@@ -341,13 +322,13 @@ export class ModalManager extends TypedEventEmitter<ModalManagerEvent, HandlerMa
         };
     }
 
-    private appendDialogAsync<C extends ComponentType>(
+    public appendDialog<C extends ComponentType>(
-        prom: Promise<C>,
+        component: C,
         props?: ComponentProps<C>,
         className?: string,
     ): IHandle<C> {
         const beforeModal = this.getCurrentModal();
-        const { modal, closeDialog, onFinishedProm } = this.buildModal<C>(prom, props, className, {});
+        const { modal, closeDialog, onFinishedProm } = this.buildModal<C>(component, props, className, {});
 
         this.modals.push(modal);
 

src/SecurityManager.ts

@@ -6,11 +6,11 @@ SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only
 Please see LICENSE files in the repository root for full details.
 */
 
+import { lazy } from "react";
 import { ICryptoCallbacks, SecretStorage } from "matrix-js-sdk/src/matrix";
 import { deriveRecoveryKeyFromPassphrase, decodeRecoveryKey } from "matrix-js-sdk/src/crypto-api";
 import { logger } from "matrix-js-sdk/src/logger";
 
-import type CreateSecretStorageDialog from "./async-components/views/dialogs/security/CreateSecretStorageDialog";
 import Modal from "./Modal";
 import { MatrixClientPeg } from "./MatrixClientPeg";
 import { _t } from "./languageHandler";
@@ -232,10 +232,8 @@ async function doAccessSecretStorage(func: () => Promise<void>, forceReset: bool
         if (createNew) {
             // This dialog calls bootstrap itself after guiding the user through
             // passphrase creation.
-            const { finished } = Modal.createDialogAsync(
+            const { finished } = Modal.createDialog(
-                import("./async-components/views/dialogs/security/CreateSecretStorageDialog") as unknown as Promise<
+                lazy(() => import("./async-components/views/dialogs/security/CreateSecretStorageDialog")),
-                    typeof CreateSecretStorageDialog
-                >,
                 {
                     forceReset,
                 },

src/async-components/views/dialogs/security/NewRecoveryMethodDialog.tsx

@@ -28,7 +28,7 @@ interface NewRecoveryMethodDialogProps {
     onFinished(): void;
 }
 
-// Export as default instead of a named export so that it can be dynamically imported with `Modal.createDialogAsync`
+// Export as default instead of a named export so that it can be dynamically imported with React lazy
 
 /**
  * Dialog to inform the user that a new recovery method has been detected.

src/async-components/views/dialogs/security/RecoveryMethodRemovedDialog.tsx

@@ -7,11 +7,11 @@ SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only
 Please see LICENSE files in the repository root for full details.
 */
 
-import React from "react";
+import React, { lazy } from "react";
 
 import dis from "../../../../dispatcher/dispatcher";
 import { _t } from "../../../../languageHandler";
-import Modal, { ComponentType } from "../../../../Modal";
+import Modal from "../../../../Modal";
 import { Action } from "../../../../dispatcher/actions";
 import BaseDialog from "../../../../components/views/dialogs/BaseDialog";
 import DialogButtons from "../../../../components/views/elements/DialogButtons";
@@ -28,8 +28,8 @@ export default class RecoveryMethodRemovedDialog extends React.PureComponent<IPr
 
     private onSetupClick = (): void => {
         this.props.onFinished();
-        Modal.createDialogAsync(
+        Modal.createDialog(
-            import("./CreateKeyBackupDialog") as unknown as Promise<ComponentType>,
+            lazy(() => import("./CreateKeyBackupDialog")),
             undefined,
             undefined,
             /* priority = */ false,

src/components/structures/MatrixChat.tsx

@@ -6,7 +6,7 @@ SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only
 Please see LICENSE files in the repository root for full details.
 */
 
-import React, { createRef } from "react";
+import React, { createRef, lazy } from "react";
 import {
     ClientEvent,
     createClient,
@@ -28,8 +28,6 @@ import { TooltipProvider } from "@vector-im/compound-web";
 // what-input helps improve keyboard accessibility
 import "what-input";
 
-import type NewRecoveryMethodDialog from "../../async-components/views/dialogs/security/NewRecoveryMethodDialog";
-import type RecoveryMethodRemovedDialog from "../../async-components/views/dialogs/security/RecoveryMethodRemovedDialog";
 import PosthogTrackers from "../../PosthogTrackers";
 import { DecryptionFailureTracker } from "../../DecryptionFailureTracker";
 import { IMatrixClientCreds, MatrixClientPeg } from "../../MatrixClientPeg";
@@ -1649,16 +1647,12 @@ export default class MatrixChat extends React.PureComponent<IProps, IState> {
             }
 
             if (haveNewVersion) {
-                Modal.createDialogAsync(
+                Modal.createDialog(
-                    import(
+                    lazy(() => import("../../async-components/views/dialogs/security/NewRecoveryMethodDialog")),
-                        "../../async-components/views/dialogs/security/NewRecoveryMethodDialog"
-                    ) as unknown as Promise<typeof NewRecoveryMethodDialog>,
                 );
             } else {
-                Modal.createDialogAsync(
+                Modal.createDialog(
-                    import(
+                    lazy(() => import("../../async-components/views/dialogs/security/RecoveryMethodRemovedDialog")),
-                        "../../async-components/views/dialogs/security/RecoveryMethodRemovedDialog"
-                    ) as unknown as Promise<typeof RecoveryMethodRemovedDialog>,
                 );
             }
         });

src/components/views/dialogs/LogoutDialog.tsx

@@ -7,12 +7,10 @@ SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only
 Please see LICENSE files in the repository root for full details.
 */
 
-import React from "react";
+import React, { lazy } from "react";
 import { logger } from "matrix-js-sdk/src/logger";
 import { MatrixClient } from "matrix-js-sdk/src/matrix";
 
-import type CreateKeyBackupDialog from "../../../async-components/views/dialogs/security/CreateKeyBackupDialog";
-import type ExportE2eKeysDialog from "../../../async-components/views/dialogs/security/ExportE2eKeysDialog";
 import Modal from "../../../Modal";
 import dis from "../../../dispatcher/dispatcher";
 import { _t } from "../../../languageHandler";
@@ -116,10 +114,8 @@ export default class LogoutDialog extends React.Component<IProps, IState> {
     }
 
     private onExportE2eKeysClicked = (): void => {
-        Modal.createDialogAsync(
+        Modal.createDialog(
-            import("../../../async-components/views/dialogs/security/ExportE2eKeysDialog") as unknown as Promise<
+            lazy(() => import("../../../async-components/views/dialogs/security/ExportE2eKeysDialog")),
-                typeof ExportE2eKeysDialog
-            >,
             {
                 matrixClient: MatrixClientPeg.safeGet(),
             },
@@ -147,10 +143,8 @@ export default class LogoutDialog extends React.Component<IProps, IState> {
                 /* static = */ true,
             );
         } else {
-            Modal.createDialogAsync(
+            Modal.createDialog(
-                import("../../../async-components/views/dialogs/security/CreateKeyBackupDialog") as unknown as Promise<
+                lazy(() => import("../../../async-components/views/dialogs/security/CreateKeyBackupDialog")),
-                    typeof CreateKeyBackupDialog
-                >,
                 undefined,
                 undefined,
                 /* priority = */ false,

src/components/views/settings/CryptographyPanel.tsx

@@ -6,11 +6,9 @@ SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only
 Please see LICENSE files in the repository root for full details.
 */
 
-import React from "react";
+import React, { lazy } from "react";
 import { logger } from "matrix-js-sdk/src/logger";
 
-import type ExportE2eKeysDialog from "../../../async-components/views/dialogs/security/ExportE2eKeysDialog";
-import type ImportE2eKeysDialog from "../../../async-components/views/dialogs/security/ImportE2eKeysDialog";
 import { _t } from "../../../languageHandler";
 import Modal from "../../../Modal";
 import AccessibleButton from "../elements/AccessibleButton";
@@ -129,19 +127,15 @@ export default class CryptographyPanel extends React.Component<IProps, IState> {
     }
 
     private onExportE2eKeysClicked = (): void => {
-        Modal.createDialogAsync(
+        Modal.createDialog(
-            import("../../../async-components/views/dialogs/security/ExportE2eKeysDialog") as unknown as Promise<
+            lazy(() => import("../../../async-components/views/dialogs/security/ExportE2eKeysDialog")),
-                typeof ExportE2eKeysDialog
-            >,
             { matrixClient: this.context },
         );
     };
 
     private onImportE2eKeysClicked = (): void => {
-        Modal.createDialogAsync(
+        Modal.createDialog(
-            import("../../../async-components/views/dialogs/security/ImportE2eKeysDialog") as unknown as Promise<
+            lazy(() => import("../../../async-components/views/dialogs/security/ImportE2eKeysDialog")),
-                typeof ImportE2eKeysDialog
-            >,
             { matrixClient: this.context },
         );
     };

src/components/views/settings/EventIndexPanel.tsx

@@ -6,7 +6,7 @@ SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only
 Please see LICENSE files in the repository root for full details.
 */
 
-import React from "react";
+import React, { lazy } from "react";
 
 import { _t } from "../../../languageHandler";
 import SdkConfig from "../../../SdkConfig";
@@ -94,14 +94,12 @@ export default class EventIndexPanel extends React.Component<{}, IState> {
     }
 
     private onManage = async (): Promise<void> => {
-        Modal.createDialogAsync(
+        Modal.createDialog(
-            // @ts-ignore: TS doesn't seem to like the type of this now that it
+            lazy(() => import("../../../async-components/views/dialogs/eventindex/ManageEventIndexDialog")),
-            // has also been converted to TS as well, but I can't figure out why...
-            import("../../../async-components/views/dialogs/eventindex/ManageEventIndexDialog"),
             {
                 onFinished: () => {},
             },
-            null,
+            undefined,
             /* priority = */ false,
             /* static = */ true,
         );

src/components/views/settings/SecureBackupPanel.tsx

@@ -7,11 +7,10 @@ SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only
 Please see LICENSE files in the repository root for full details.
 */
 
-import React, { ReactNode } from "react";
+import React, { lazy, ReactNode } from "react";
 import { CryptoEvent, BackupTrustInfo, KeyBackupInfo } from "matrix-js-sdk/src/crypto-api";
 import { logger } from "matrix-js-sdk/src/logger";
 
-import type CreateKeyBackupDialog from "../../../async-components/views/dialogs/security/CreateKeyBackupDialog";
 import { MatrixClientPeg } from "../../../MatrixClientPeg";
 import { _t } from "../../../languageHandler";
 import Modal from "../../../Modal";
@@ -170,10 +169,8 @@ export default class SecureBackupPanel extends React.PureComponent<{}, IState> {
     }
 
     private startNewBackup = (): void => {
-        Modal.createDialogAsync(
+        Modal.createDialog(
-            import("../../../async-components/views/dialogs/security/CreateKeyBackupDialog") as unknown as Promise<
+            lazy(() => import("../../../async-components/views/dialogs/security/CreateKeyBackupDialog")),
-                typeof CreateKeyBackupDialog
-            >,
             {
                 onFinished: () => {
                     this.loadBackupStatus();

test/unit-tests/SecurityManager-test.ts

@@ -11,6 +11,13 @@ import { CryptoApi } from "matrix-js-sdk/src/crypto-api";
 
 import { accessSecretStorage } from "../../src/SecurityManager";
 import { filterConsole, stubClient } from "../test-utils";
+import Modal from "../../src/Modal.tsx";
+
+jest.mock("react", () => {
+    const React = jest.requireActual("react");
+    React.lazy = (children: any) => children(); // stub out lazy for dialog test
+    return React;
+});
 
 describe("SecurityManager", () => {
     describe("accessSecretStorage", () => {
@@ -50,5 +57,21 @@ describe("SecurityManager", () => {
                 }).rejects.toThrow("End-to-end encryption is disabled - unable to access secret storage");
             });
         });
+
+        it("should show CreateSecretStorageDialog if forceReset=true", async () => {
+            jest.mock("../../src/async-components/views/dialogs/security/CreateSecretStorageDialog", () => ({
+                __test: true,
+                __esModule: true,
+                default: () => jest.fn(),
+            }));
+            const spy = jest.spyOn(Modal, "createDialog");
+            stubClient();
+
+            const func = jest.fn();
+            accessSecretStorage(func, true);
+
+            expect(spy).toHaveBeenCalledTimes(1);
+            await expect(spy.mock.lastCall![0]).resolves.toEqual(expect.objectContaining({ __test: true }));
+        });
     });
 });

test/unit-tests/async-components/dialogs/security/RecoveryMethodRemovedDialog-test.tsx

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2024 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only
+ * Please see LICENSE files in the repository root for full details.
+ */
+
+import React from "react";
+import { render, screen, fireEvent, waitFor } from "jest-matrix-react";
+
+import RecoveryMethodRemovedDialog from "../../../../../src/async-components/views/dialogs/security/RecoveryMethodRemovedDialog";
+import Modal from "../../../../../src/Modal.tsx";
+
+describe("<RecoveryMethodRemovedDialog />", () => {
+    afterEach(() => {
+        jest.restoreAllMocks();
+    });
+
+    it("should open CreateKeyBackupDialog on primary action click", async () => {
+        const onFinished = jest.fn();
+        const spy = jest.spyOn(Modal, "createDialog");
+        jest.mock("../../../../../src/async-components/views/dialogs/security/CreateKeyBackupDialog", () => ({
+            __test: true,
+            __esModule: true,
+            default: () => <span>mocked dialog</span>,
+        }));
+
+        render(<RecoveryMethodRemovedDialog onFinished={onFinished} />);
+        fireEvent.click(screen.getByRole("button", { name: "Set up Secure Messages" }));
+        await waitFor(() => expect(spy).toHaveBeenCalledTimes(1));
+        expect((spy.mock.lastCall![0] as any)._payload._result).toEqual(expect.objectContaining({ __test: true }));
+    });
+});

test/unit-tests/components/structures/MatrixChat-test.tsx

@@ -149,6 +149,7 @@ describe("<MatrixChat />", () => {
         isRoomEncrypted: jest.fn(),
         logout: jest.fn(),
         getDeviceId: jest.fn(),
+        getKeyBackupVersion: jest.fn().mockResolvedValue(null),
     });
     let mockClient: Mocked<MatrixClient>;
     const serverConfig = {
@@ -1515,7 +1516,7 @@ describe("<MatrixChat />", () => {
 
     describe("when key backup failed", () => {
         it("should show the new recovery method dialog", async () => {
-            const spy = jest.spyOn(Modal, "createDialogAsync");
+            const spy = jest.spyOn(Modal, "createDialog");
             jest.mock("../../../../src/async-components/views/dialogs/security/NewRecoveryMethodDialog", () => ({
                 __test: true,
                 __esModule: true,
@@ -1530,7 +1531,25 @@ describe("<MatrixChat />", () => {
             await flushPromises();
             mockClient.emit(CryptoEvent.KeyBackupFailed, "error code");
             await waitFor(() => expect(spy).toHaveBeenCalledTimes(1));
-            expect(await spy.mock.lastCall![0]).toEqual(expect.objectContaining({ __test: true }));
+            expect((spy.mock.lastCall![0] as any)._payload._result).toEqual(expect.objectContaining({ __test: true }));
+        });
+
+        it("should show the recovery method removed dialog", async () => {
+            const spy = jest.spyOn(Modal, "createDialog");
+            jest.mock("../../../../src/async-components/views/dialogs/security/RecoveryMethodRemovedDialog", () => ({
+                __test: true,
+                __esModule: true,
+                default: () => <span>mocked dialog</span>,
+            }));
+
+            getComponent({});
+            defaultDispatcher.dispatch({
+                action: "will_start_client",
+            });
+            await flushPromises();
+            mockClient.emit(CryptoEvent.KeyBackupFailed, "error code");
+            await waitFor(() => expect(spy).toHaveBeenCalledTimes(1));
+            expect((spy.mock.lastCall![0] as any)._payload._result).toEqual(expect.objectContaining({ __test: true }));
         });
     });
 });

test/unit-tests/components/views/dialogs/LogoutDialog-test.tsx

@@ -10,7 +10,7 @@ import React from "react";
 import { mocked, MockedObject } from "jest-mock";
 import { MatrixClient } from "matrix-js-sdk/src/matrix";
 import { CryptoApi, KeyBackupInfo } from "matrix-js-sdk/src/crypto-api";
-import { render, RenderResult } from "jest-matrix-react";
+import { fireEvent, render, RenderResult, screen } from "jest-matrix-react";
 
 import { filterConsole, getMockClientWithEventEmitter, mockClientMethodsCrypto } from "../../../../test-utils";
 import LogoutDialog from "../../../../../src/components/views/dialogs/LogoutDialog";
@@ -61,6 +61,9 @@ describe("LogoutDialog", () => {
         const rendered = renderComponent();
         await rendered.findByText("Start using Key Backup");
         expect(rendered.container).toMatchSnapshot();
+
+        fireEvent.click(await screen.findByRole("button", { name: "Manually export keys" }));
+        await expect(screen.findByRole("heading", { name: "Export room keys" })).resolves.toBeInTheDocument();
     });
 
     describe("when there is an error fetching backups", () => {

test/unit-tests/components/views/settings/AddRemoveThreepids-test.tsx

@@ -7,7 +7,7 @@ Please see LICENSE files in the repository root for full details.
 */
 
 import { render, screen, waitFor } from "jest-matrix-react";
-import { MatrixClient, ThreepidMedium } from "matrix-js-sdk/src/matrix";
+import { MatrixClient, MatrixError, ThreepidMedium } from "matrix-js-sdk/src/matrix";
 import React from "react";
 import userEvent from "@testing-library/user-event";
 import { mocked } from "jest-mock";
@@ -16,6 +16,7 @@ import { AddRemoveThreepids } from "../../../../../src/components/views/settings
 import { clearAllModals, stubClient } from "../../../../test-utils";
 import MatrixClientContext from "../../../../../src/contexts/MatrixClientContext";
 import Modal from "../../../../../src/Modal";
+import InteractiveAuthDialog from "../../../../../src/components/views/dialogs/InteractiveAuthDialog.tsx";
 
 const MOCK_IDENTITY_ACCESS_TOKEN = "mock_identity_access_token";
 const mockGetAccessToken = jest.fn().mockResolvedValue(MOCK_IDENTITY_ACCESS_TOKEN);
@@ -222,13 +223,13 @@ describe("AddRemoveThreepids", () => {
 
         const continueButton = await screen.findByRole("button", { name: /Continue/ });
 
-        await expect(continueButton).toHaveAttribute("aria-disabled", "true");
+        expect(continueButton).toHaveAttribute("aria-disabled", "true");
 
         await expect(
-            await screen.findByText(
+            screen.findByText(
                 `A text message has been sent to +${PHONE1.address}. Please enter the verification code it contains.`,
             ),
-        ).toBeInTheDocument();
+        ).resolves.toBeInTheDocument();
 
         expect(client.requestAdd3pidMsisdnToken).toHaveBeenCalledWith(
             "GB",
@@ -481,4 +482,118 @@ describe("AddRemoveThreepids", () => {
         expect(client.unbindThreePid).toHaveBeenCalledWith(ThreepidMedium.Phone, PHONE1.address);
         expect(onChangeFn).toHaveBeenCalled();
     });
+
+    it("should show UIA dialog when necessary for adding email", async () => {
+        const onChangeFn = jest.fn();
+        const createDialogFn = jest.spyOn(Modal, "createDialog");
+        mocked(client.requestAdd3pidEmailToken).mockResolvedValue({ sid: "1" });
+
+        render(
+            <AddRemoveThreepids
+                mode="hs"
+                medium={ThreepidMedium.Email}
+                threepids={[]}
+                isLoading={false}
+                onChange={onChangeFn}
+            />,
+            {
+                wrapper: clientProviderWrapper,
+            },
+        );
+
+        const input = screen.getByRole("textbox", { name: "Email Address" });
+        await userEvent.type(input, EMAIL1.address);
+        const addButton = screen.getByRole("button", { name: "Add" });
+        await userEvent.click(addButton);
+
+        const continueButton = screen.getByRole("button", { name: "Continue" });
+
+        expect(continueButton).toBeEnabled();
+
+        mocked(client).addThreePidOnly.mockRejectedValueOnce(
+            new MatrixError({ errcode: "M_UNAUTHORIZED", flows: [{ stages: [] }] }, 401),
+        );
+
+        await userEvent.click(continueButton);
+
+        expect(createDialogFn).toHaveBeenCalledWith(
+            InteractiveAuthDialog,
+            expect.objectContaining({
+                title: "Add Email Address",
+                makeRequest: expect.any(Function),
+            }),
+        );
+    });
+
+    it("should show UIA dialog when necessary for adding msisdn", async () => {
+        const onChangeFn = jest.fn();
+        const createDialogFn = jest.spyOn(Modal, "createDialog");
+        mocked(client.requestAdd3pidMsisdnToken).mockResolvedValue({
+            sid: "1",
+            msisdn: PHONE1.address,
+            intl_fmt: PHONE1.address,
+            success: true,
+            submit_url: "https://some-url",
+        });
+
+        render(
+            <AddRemoveThreepids
+                mode="hs"
+                medium={ThreepidMedium.Phone}
+                threepids={[]}
+                isLoading={false}
+                onChange={onChangeFn}
+            />,
+            {
+                wrapper: clientProviderWrapper,
+            },
+        );
+
+        const countryDropdown = screen.getByRole("button", { name: /Country Dropdown/ });
+        await userEvent.click(countryDropdown);
+        const gbOption = screen.getByRole("option", { name: "🇬🇧 United Kingdom (+44)" });
+        await userEvent.click(gbOption);
+
+        const input = screen.getByRole("textbox", { name: "Phone Number" });
+        await userEvent.type(input, PHONE1_LOCALNUM);
+
+        const addButton = screen.getByRole("button", { name: "Add" });
+        await userEvent.click(addButton);
+
+        const continueButton = screen.getByRole("button", { name: "Continue" });
+
+        expect(continueButton).toHaveAttribute("aria-disabled", "true");
+
+        await expect(
+            screen.findByText(
+                `A text message has been sent to +${PHONE1.address}. Please enter the verification code it contains.`,
+            ),
+        ).resolves.toBeInTheDocument();
+
+        expect(client.requestAdd3pidMsisdnToken).toHaveBeenCalledWith(
+            "GB",
+            PHONE1_LOCALNUM,
+            client.generateClientSecret(),
+            1,
+        );
+
+        const verificationInput = screen.getByRole("textbox", { name: "Verification code" });
+        await userEvent.type(verificationInput, "123456");
+
+        expect(continueButton).not.toHaveAttribute("aria-disabled", "true");
+
+        mocked(client).addThreePidOnly.mockRejectedValueOnce(
+            new MatrixError({ errcode: "M_UNAUTHORIZED", flows: [{ stages: [] }] }, 401),
+        );
+
+        await userEvent.click(continueButton);
+
+        expect(createDialogFn).toHaveBeenCalledWith(
+            InteractiveAuthDialog,
+            expect.objectContaining({
+                title: "Add Phone Number",
+                makeRequest: expect.any(Function),
+            }),
+        );
+    });
 });

test/unit-tests/components/views/settings/CryptographyPanel-test.tsx

@@ -7,7 +7,7 @@ Please see LICENSE files in the repository root for full details.
 */
 
 import React from "react";
-import { render, waitFor } from "jest-matrix-react";
+import { render, waitFor, screen, fireEvent } from "jest-matrix-react";
 import { MatrixClient } from "matrix-js-sdk/src/matrix";
 import { mocked } from "jest-mock";
 
@@ -64,4 +64,34 @@ describe("CryptographyPanel", () => {
         // Then "not supported key
         await waitFor(() => expect(codes[1].innerHTML).toEqual("<strong>&lt;not supported&gt;</strong>"));
     });
+
+    it("should open the export e2e keys dialog on click", async () => {
+        const sessionId = "ABCDEFGHIJ";
+        const sessionKey = "AbCDeFghIJK7L/m4nOPqRSTUVW4xyzaBCDef6gHIJkl";
+
+        TestUtils.stubClient();
+        const client: MatrixClient = MatrixClientPeg.safeGet();
+        client.deviceId = sessionId;
+
+        mocked(client.getCrypto()!.getOwnDeviceKeys).mockResolvedValue({ ed25519: sessionKey, curve25519: "1234" });
+
+        render(<CryptographyPanel />, withClientContextRenderOptions(client));
+        fireEvent.click(await screen.findByRole("button", { name: "Export E2E room keys" }));
+        await expect(screen.findByRole("heading", { name: "Export room keys" })).resolves.toBeInTheDocument();
+    });
+
+    it("should open the import e2e keys dialog on click", async () => {
+        const sessionId = "ABCDEFGHIJ";
+        const sessionKey = "AbCDeFghIJK7L/m4nOPqRSTUVW4xyzaBCDef6gHIJkl";
+
+        TestUtils.stubClient();
+        const client: MatrixClient = MatrixClientPeg.safeGet();
+        client.deviceId = sessionId;
+
+        mocked(client.getCrypto()!.getOwnDeviceKeys).mockResolvedValue({ ed25519: sessionKey, curve25519: "1234" });
+
+        render(<CryptographyPanel />, withClientContextRenderOptions(client));
+        fireEvent.click(await screen.findByRole("button", { name: "Import E2E room keys" }));
+        await expect(screen.findByRole("heading", { name: "Import room keys" })).resolves.toBeInTheDocument();
+    });
 });