#3509592 Attempted fixes for test flake

Closes #3509592

ui/src/hooks/useSyncIframeHeightToContent.ts

@@ -39,43 +39,58 @@ function useSyncIframeHeightToContent(
 
   useEffect(() => {
     if (iframe) {
-      const iframeContentDoc = iframe.contentDocument;
+      const handleLoad = () => {
+        const iframeContentDoc = iframe.contentDocument;
 
-      if (iframeContentDoc) {
-        const iframeHTML = iframeContentDoc.documentElement;
+        if (iframeContentDoc) {
+          const iframeHTML = iframeContentDoc.documentElement;
 
-        iframeHTML.style.overflow = 'hidden';
-        // Set up a MutationObserver to watch for changes in the content of the iframe
-        mutationObserverRef.current = new MutationObserver(resizeIframe);
-        mutationObserverRef.current.observe(iframeHTML, {
-          attributes: true,
-          childList: true,
-          subtree: true,
-        });
-        resizeObserverRef.current = new ResizeObserver(resizeIframe);
-        resizeObserverRef.current.observe(iframeHTML);
+          iframeHTML.style.overflow = 'hidden';
+          // Set up a MutationObserver to watch for changes in the content of the iframe
+          mutationObserverRef.current = new MutationObserver(resizeIframe);
+          mutationObserverRef.current.observe(iframeHTML, {
+            attributes: true,
+            childList: true,
+            subtree: true,
+          });
+          resizeObserverRef.current = new ResizeObserver(resizeIframe);
+          resizeObserverRef.current.observe(iframeHTML);
 
-        // Apply a max-height to elements with vh units in their height - otherwise an infinite loop can occur where a component's
-        // height is based on the height of the iFrame and the iFrame's height is based on that component leading
-        // to an ever-increasing iFrame height!
-        const elements: NodeListOf<HTMLElement> =
-          iframeHTML.querySelectorAll('*');
-        elements.forEach((element: HTMLElement) => {
-          if (element.style.height.endsWith('vh')) {
-            const vhValue = parseFloat(element.style.height);
-            const newHeight = (vhValue / 100) * height;
-            element.style.maxHeight = newHeight + 'px';
-            resizeIframe();
-          }
-        });
+          // Apply a max-height to elements with vh units in their height - otherwise an infinite loop can occur where a component's
+          // height is based on the height of the iFrame and the iFrame's height is based on that component leading
+          // to an ever-increasing iFrame height!
+          const elements: NodeListOf<HTMLElement> =
+            iframeHTML.querySelectorAll('*');
+          elements.forEach((element) => {
+            if (
+              element instanceof HTMLElement &&
+              element.style.height.endsWith('vh')
+            ) {
+              const vhValue = parseFloat(element.style.height);
+              const newHeight = (vhValue / 100) * height;
+              element.style.maxHeight = newHeight + 'px';
+              resizeIframe();
+            }
+          });
 
-        resizeIframe();
+          resizeIframe();
+        }
+      };
+
+      // Assign the load event listener
+      iframe.addEventListener('load', handleLoad);
+
+      // Check if the iFrame is already loaded
+      if (iframe.contentDocument?.readyState === 'complete') {
+        handleLoad();
       }
+
+      return () => {
+        iframe.removeEventListener('load', handleLoad);
+        mutationObserverRef.current?.disconnect();
+        resizeObserverRef.current?.disconnect();
+      };
     }
-    return () => {
-      mutationObserverRef.current?.disconnect();
-      resizeObserverRef.current?.disconnect();
-    };
   }, [iframe, height, resizeIframe]);
 }