From 72fd50c3ecffe93704ea8eae15252b10d13f403c Mon Sep 17 00:00:00 2001
From: Daphne Li-Chen <104820175+dlichen@users.noreply.github.com>
Date: Fri, 27 May 2022 09:20:37 -0700
Subject: [PATCH] [core] feat(Dialog): add containerRef prop (#5314)

Co-authored-by: Adi Dahiya <adi.dahiya14@gmail.com>
Co-authored-by: Adi Dahiya <adahiya@palantir.com>
---
 .../core/src/components/dialog/dialog.tsx     |  9 +++-
 packages/core/test/dialog/dialogTests.tsx     | 45 ++++++++++---------
 2 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/packages/core/src/components/dialog/dialog.tsx b/packages/core/src/components/dialog/dialog.tsx
index 5e5acda922..fe374d7265 100644
--- a/packages/core/src/components/dialog/dialog.tsx
+++ b/packages/core/src/components/dialog/dialog.tsx
@@ -17,7 +17,7 @@
 import classNames from "classnames";
 import * as React from "react";
 
-import { AbstractPureComponent2, Classes } from "../../common";
+import { AbstractPureComponent2, Classes, IRef } from "../../common";
 import * as Errors from "../../common/errors";
 import { DISPLAYNAME_PREFIX, MaybeElement, Props } from "../../common/props";
 import { uniqueId } from "../../common/utils";
@@ -80,6 +80,11 @@ export interface IDialogProps extends OverlayableProps, IBackdropProps, Props {
      */
     transitionName?: string;
 
+    /**
+     * Ref supplied to the `Classes.DIALOG_CONTAINER` element.
+     */
+    containerRef?: IRef<HTMLDivElement>;
+
     /**
      * ID of the element that contains title or label text for this dialog.
      *
@@ -114,7 +119,7 @@ export class Dialog extends AbstractPureComponent2<DialogProps> {
     public render() {
         return (
             <Overlay {...this.props} className={Classes.OVERLAY_SCROLL_CONTAINER} hasBackdrop={true}>
-                <div className={Classes.DIALOG_CONTAINER}>
+                <div className={Classes.DIALOG_CONTAINER} ref={this.props.containerRef}>
                     <div
                         className={classNames(Classes.DIALOG, this.props.className)}
                         role="dialog"
diff --git a/packages/core/test/dialog/dialogTests.tsx b/packages/core/test/dialog/dialogTests.tsx
index ff0da61135..42f12d6f20 100644
--- a/packages/core/test/dialog/dialogTests.tsx
+++ b/packages/core/test/dialog/dialogTests.tsx
@@ -19,7 +19,7 @@ import { mount } from "enzyme";
 import * as React from "react";
 import { spy } from "sinon";
 
-import { Button, Classes, Dialog, H4, Icon, IconSize } from "../../src";
+import { Button, Classes, Dialog, DialogProps, H4, Icon, IconSize } from "../../src";
 import * as Keys from "../../src/common/keys";
 
 describe("<Dialog>", () => {
@@ -147,54 +147,55 @@ describe("<Dialog>", () => {
     });
 
     describe("accessibility features", () => {
-        const mountDialog = (className: string) => {
+        const mountDialog = (props: Partial<DialogProps>) => {
             return mount(
-                <Dialog
-                    className={className}
-                    isOpen={true}
-                    usePortal={false}
-                    aria-labelledby="dialog-title"
-                    aria-describedby="dialog-description"
-                >
+                <Dialog isOpen={true} usePortal={false} {...props}>
                     {createDialogContents()}
                 </Dialog>,
             );
         };
 
         it("renders with role={dialog}", () => {
-            const dialog = mountDialog("check-role");
+            const dialog = mountDialog({ className: "check-role" });
             assert.equal(dialog.find(`.check-role`).hostNodes().prop("role"), "dialog", "missing dialog role!!");
         });
 
         it("renders with provided aria-labelledby and aria-described by from props", () => {
-            const dialog = mountDialog("renders-with-props");
+            const dialog = mountDialog({
+                "aria-describedby": "dialog-description",
+                "aria-labelledby": "dialog-title",
+                className: "renders-with-props",
+            });
             const dialogElement = dialog.find(`.renders-with-props`).hostNodes();
             assert.equal(dialogElement.prop("aria-labelledby"), "dialog-title");
             assert.equal(dialogElement.prop("aria-describedby"), "dialog-description");
         });
 
         it("uses title as default aria-labelledby", () => {
-            const dialog = mount(
-                <Dialog className="default-title" isOpen={true} usePortal={false} title="Title by props">
-                    {createDialogContents()}
-                </Dialog>,
-            );
+            const dialog = mountDialog({ className: "default-title", title: "Title by props" });
             // test existence here because id is generated
             assert.exists(dialog.find(".default-title").hostNodes().prop("aria-labelledby"));
         });
 
         it("does not apply default aria-labelledby if no title", () => {
-            const dialog = mount(
-                <Dialog className={"no-default-if-no-title"} isOpen={true} usePortal={false}>
-                    {createDialogContents()}
-                </Dialog>,
-            );
+            const dialog = mountDialog({ className: "no-default-if-no-title" });
             // test existence here because id is generated
             assert.notExists(dialog.find(".no-default-if-no-title").hostNodes().prop("aria-labelledby"));
         });
+
+        it("supports ref objects attached to container", done => {
+            const containerRef = React.createRef<HTMLDivElement>();
+            mountDialog({ containerRef });
+
+            // wait for the whole lifecycle to run
+            setTimeout(() => {
+                assert.isTrue(containerRef.current?.classList.contains(Classes.DIALOG_CONTAINER));
+                done();
+            }, 0);
+        });
     });
 
-    // everything else about Dialog is tested by Overlay
+    // N.B. everything else about Dialog is tested by Overlay
 
     function createDialogContents(): JSX.Element[] {
         return [