From 7f9c74cd8cc747a6a578825ab1ec1a74a53457c6 Mon Sep 17 00:00:00 2001
From: Bartosz Prusinowski <bartosz@interactivethings.com>
Date: Thu, 5 Dec 2024 16:09:24 +0100
Subject: [PATCH] fix: Make modifyNode async and await animation frame where
 needed

---
 app/components/chart-shared.tsx |  4 ++--
 app/utils/animation-frame.ts    |  3 +++
 app/utils/use-screenshot.ts     | 22 ++++++++++++++++------
 3 files changed, 21 insertions(+), 8 deletions(-)
 create mode 100644 app/utils/animation-frame.ts

diff --git a/app/components/chart-shared.tsx b/app/components/chart-shared.tsx
index 0730e94a3..ea2c74b15 100644
--- a/app/components/chart-shared.tsx
+++ b/app/components/chart-shared.tsx
@@ -55,6 +55,7 @@ import { useDataCubesMetadataQuery } from "@/graphql/hooks";
 import { getChartIcon } from "@/icons";
 import SvgIcMore from "@/icons/components/IcMore";
 import { useLocale } from "@/src";
+import { animationFrame } from "@/utils/animation-frame";
 import { createChartId } from "@/utils/create-chart-id";
 import {
   DISABLE_SCREENSHOT_ATTR,
@@ -459,8 +460,7 @@ const useModifyNode = () => {
             createdWith={t({ id: "metadata.link.created.with" })}
           />
         );
-        // Wait for the component to render before taking the screenshot.
-        await new Promise((resolve) => setTimeout(resolve, 0));
+        await animationFrame();
       }
 
       // Remove some elements that should not be included in the screenshot.
diff --git a/app/utils/animation-frame.ts b/app/utils/animation-frame.ts
new file mode 100644
index 000000000..475246540
--- /dev/null
+++ b/app/utils/animation-frame.ts
@@ -0,0 +1,3 @@
+export const animationFrame = () => {
+  return new Promise((resolve) => requestAnimationFrame(resolve));
+};
diff --git a/app/utils/use-screenshot.ts b/app/utils/use-screenshot.ts
index d0e997b4e..827040911 100644
--- a/app/utils/use-screenshot.ts
+++ b/app/utils/use-screenshot.ts
@@ -3,13 +3,18 @@ import { toPng, toSvg } from "html-to-image";
 import { addMetadata } from "meta-png";
 import { useCallback, useState } from "react";
 
+import { animationFrame } from "@/utils/animation-frame";
+
 type ScreenshotFileFormat = "png" | "svg";
 
 export type UseScreenshotProps = {
   type: ScreenshotFileFormat;
   screenshotName: string;
   screenshotNode?: HTMLElement | null;
-  modifyNode?: (clonedNode: HTMLElement, originalNode: HTMLElement) => void;
+  modifyNode?: (
+    clonedNode: HTMLElement,
+    originalNode: HTMLElement
+  ) => Promise<void>;
   pngMetadata?: { key: PNG_METADATA_KEY; value: string }[];
 };
 
@@ -22,11 +27,11 @@ export const useScreenshot = ({
 }: UseScreenshotProps) => {
   const [loading, setLoading] = useState(false);
   const modifyNode = useCallback(
-    (clonedNode: HTMLElement, originalNode: HTMLElement) => {
+    async (clonedNode: HTMLElement, originalNode: HTMLElement) => {
       removeDisabledElements(clonedNode);
 
       if (_modifyNode) {
-        _modifyNode(clonedNode, originalNode);
+        await _modifyNode(clonedNode, originalNode);
       }
     },
     [_modifyNode]
@@ -81,7 +86,10 @@ const makeScreenshot = async ({
   type: "png" | "svg";
   name: string;
   node: HTMLElement;
-  modifyNode?: (clonedNode: HTMLElement, originalNode: HTMLElement) => void;
+  modifyNode?: (
+    clonedNode: HTMLElement,
+    originalNode: HTMLElement
+  ) => Promise<void>;
   pngMetadata?: { key: PNG_METADATA_KEY; value: string }[];
 }) => {
   const isUsingSafari =
@@ -108,10 +116,12 @@ const makeScreenshot = async ({
     // manually copy it.
     const ctx = clonedCanvasNode.getContext("2d") as CanvasRenderingContext2D;
     ctx.drawImage(canvasNode, 0, 0);
+    await animationFrame();
   }
 
-  const mountedClonedNode = wrapperNode.appendChild(clonedNode);
-  modifyNode?.(mountedClonedNode, node);
+  await modifyNode?.(clonedNode, node);
+  wrapperNode.appendChild(clonedNode);
+  await animationFrame();
 
   // There's a bug with embedding the fonts in Safari, which appears only when
   // downloading the image for the first time. On subsequent downloads, the