From 98adb14a8ffff43baf47e9ef7de2ae96791056e0 Mon Sep 17 00:00:00 2001
From: Jakub Szewczyk <jakub.szewczykk@gmail.com>
Date: Mon, 28 Mar 2022 18:02:10 +0200
Subject: [PATCH] Fix useRewards hook that didn't work with async components
 (#84)

Co-authored-by: jakub.szewczyk <jakub.szewczyk@tivix.com>
---
 package.json             |  2 +-
 src/functions/helpers.ts | 10 +++++++++
 src/hooks/useReward.ts   | 46 +++++++++++++++++-----------------------
 3 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/package.json b/package.json
index acdaed8..e740d54 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "react-rewards",
-  "version": "2.0.2",
+  "version": "2.0.3",
   "description": "This package lets you easily add micro-interactions to your app and reward users with the rain of confetti, emoji or balloons.",
   "author": "thedevelobear",
   "license": "MIT",
diff --git a/src/functions/helpers.ts b/src/functions/helpers.ts
index 50a5998..90aca65 100644
--- a/src/functions/helpers.ts
+++ b/src/functions/helpers.ts
@@ -64,3 +64,13 @@ export const generatePhysics = (
     differentiator,
   };
 };
+
+export const getContainerById = (id: string) => {
+  const container = document.getElementById(id);
+  if (!container) {
+    console.error(
+      `Element with an ID of ${id} could not be found. Please provide a valid ID.`
+    );
+  }
+  return container;
+};
diff --git a/src/hooks/useReward.ts b/src/hooks/useReward.ts
index 87bd9ae..31e3494 100644
--- a/src/hooks/useReward.ts
+++ b/src/hooks/useReward.ts
@@ -1,34 +1,13 @@
-import { useLayoutEffect, useState } from 'react';
+import { useState } from 'react';
 import { confetti } from '../components/Confetti/Confetti';
 import { emoji } from '../components/Emoji/Emoji';
 import { balloons } from '../components/Balloons/Balloons';
 import { UseRewardType } from './useReward.types';
+import { getContainerById } from '../functions/helpers';
 
 export const useReward: UseRewardType = (id, type, config) => {
-  const [container, setContainer] = useState<Element>();
   const [isAnimating, setIsAnimating] = useState<boolean>(false);
 
-  const elementNotFoundMessage = `Element with an ID of ${id} could not be found. Please provide a valid ID.`;
-  const typeNotFoundMessage = `${type} is not a valid react-rewards type.`;
-
-  useLayoutEffect(() => {
-    const foundContainer = document.getElementById(id);
-    if (foundContainer) {
-      setContainer(foundContainer);
-    } else {
-      console.error(elementNotFoundMessage);
-    }
-  }, [id, elementNotFoundMessage]);
-
-  if (!container) {
-    return {
-      reward: () => {
-        console.error(elementNotFoundMessage);
-      },
-      isAnimating: false,
-    };
-  }
-
   const internalAnimatingCallback = () => {
     setIsAnimating(false);
   };
@@ -37,27 +16,40 @@ export const useReward: UseRewardType = (id, type, config) => {
   switch (type) {
     case 'confetti': {
       reward = () => {
+        const foundContainer = getContainerById(id);
+
+        if (!foundContainer) return;
+
         setIsAnimating(true);
-        confetti(container, internalAnimatingCallback, config);
+        confetti(foundContainer, internalAnimatingCallback, config);
       };
       break;
     }
     case 'emoji': {
       reward = () => {
+        const foundContainer = getContainerById(id);
+
+        if (!foundContainer) return;
+
         setIsAnimating(true);
-        emoji(container, internalAnimatingCallback, config);
+        emoji(foundContainer, internalAnimatingCallback, config);
       };
       break;
     }
     case 'balloons': {
       reward = () => {
+        const foundContainer = getContainerById(id);
+
+        if (!foundContainer) return;
+
         setIsAnimating(true);
-        balloons(container, internalAnimatingCallback, config);
+        balloons(foundContainer, internalAnimatingCallback, config);
       };
       break;
     }
     default: {
-      reward = () => console.error(typeNotFoundMessage);
+      reward = () =>
+        console.error(`${type} is not a valid react-rewards type.`);
     }
   }
   return { reward, isAnimating };