From 498960312b6d0f744a773cfebac739a77553b3d5 Mon Sep 17 00:00:00 2001
From: Alexandre Philibeaux <aphilibeaux@scaleway.com>
Date: Thu, 25 Apr 2024 13:56:00 +0000
Subject: [PATCH] fix(cookie-consent): save consent doesn't trigger load of
 segment

---
 .changeset/odd-peas-suffer.md                 |  5 ++++
 .../CookieConsentProvider.tsx                 | 29 ++++++++++++-------
 .../CookieConsentProvider/__tests__/index.tsx | 13 ++++++++-
 3 files changed, 36 insertions(+), 11 deletions(-)
 create mode 100644 .changeset/odd-peas-suffer.md

diff --git a/.changeset/odd-peas-suffer.md b/.changeset/odd-peas-suffer.md
new file mode 100644
index 000000000..91cf0f816
--- /dev/null
+++ b/.changeset/odd-peas-suffer.md
@@ -0,0 +1,5 @@
+---
+"@scaleway/cookie-consent": patch
+---
+
+Fix issue with initialization on save content doens't load segment on first render.
diff --git a/packages/cookie-consent/src/CookieConsentProvider/CookieConsentProvider.tsx b/packages/cookie-consent/src/CookieConsentProvider/CookieConsentProvider.tsx
index 544a203ae..04c6513f7 100644
--- a/packages/cookie-consent/src/CookieConsentProvider/CookieConsentProvider.tsx
+++ b/packages/cookie-consent/src/CookieConsentProvider/CookieConsentProvider.tsx
@@ -52,8 +52,6 @@ export const useCookieConsent = (): Context => {
   return context
 }
 
-const getCookies = () => cookie.parse(document.cookie)
-
 export const CookieConsentProvider = ({
   children,
   isConsentRequired,
@@ -73,6 +71,9 @@ export const CookieConsentProvider = ({
   cookiesOptions?: CookieSerializeOptions
 }>) => {
   const [needConsent, setNeedsConsent] = useState(false)
+  const [cookies, setCookies] = useState<Record<string, string>>(
+    cookie.parse(document.cookie),
+  )
 
   const {
     integrations: segmentIntegrations,
@@ -113,10 +114,10 @@ export const CookieConsentProvider = ({
     // to false after receiving segment answer and flicker the UI
     setNeedsConsent(
       isConsentRequired &&
-        getCookies()[HASH_COOKIE] !== integrationsHash.toString() &&
+        cookies[HASH_COOKIE] !== integrationsHash.toString() &&
         segmentIntegrations !== undefined,
     )
-  }, [isConsentRequired, integrationsHash, segmentIntegrations])
+  }, [isConsentRequired, integrationsHash, segmentIntegrations, cookies])
 
   // We store unique categories names in an array
   const categories = useMemo(
@@ -135,13 +136,14 @@ export const CookieConsentProvider = ({
       categories.reduce<Partial<Consent>>(
         (acc, category) => ({
           ...acc,
-          [category]: isConsentRequired
-            ? getCookies()[`${cookiePrefix}_${category}`] === 'true'
-            : true,
+          [category]:
+            isConsentRequired || needConsent
+              ? cookies[`${cookiePrefix}_${category}`] === 'true'
+              : true,
         }),
         {},
       ),
-    [isConsentRequired, categories, cookiePrefix],
+    [isConsentRequired, categories, cookiePrefix, needConsent, cookies],
   )
 
   const saveConsent = useCallback(
@@ -163,7 +165,7 @@ export const CookieConsentProvider = ({
           })
         } else {
           document.cookie = cookie.serialize(
-            `${cookiePrefix}_${consentCategoryName}`,
+            cookieName,
             consentValue.toString(),
             {
               ...cookiesOptions,
@@ -174,6 +176,10 @@ export const CookieConsentProvider = ({
             },
           )
         }
+        setCookies(prevCookies => ({
+          ...prevCookies,
+          [cookieName]: consentValue ? 'true' : 'false',
+        }))
       }
       // We set the hash cookie to the current consented integrations
       document.cookie = cookie.serialize(
@@ -185,6 +191,10 @@ export const CookieConsentProvider = ({
           maxAge: consentAdvertisingMaxAge,
         },
       )
+      setCookies(prevCookies => ({
+        ...prevCookies,
+        [HASH_COOKIE]: integrationsHash.toString(),
+      }))
       setNeedsConsent(false)
     },
     [
@@ -206,7 +216,6 @@ export const CookieConsentProvider = ({
         : true,
     [isConsentRequired, segmentIntegrations, cookieConsent, needConsent],
   )
-
   // 'All': false tells Segment not to send data to any Destinations by default, unless they’re explicitly listed as true in the next lines.
   // In this case we should not have any integration, so we protect the user. Maybe unecessary as we always set true of false for an integration.
   const segmentEnabledIntegrations = useMemo(
diff --git a/packages/cookie-consent/src/CookieConsentProvider/__tests__/index.tsx b/packages/cookie-consent/src/CookieConsentProvider/__tests__/index.tsx
index c9699ab49..416d40d1c 100644
--- a/packages/cookie-consent/src/CookieConsentProvider/__tests__/index.tsx
+++ b/packages/cookie-consent/src/CookieConsentProvider/__tests__/index.tsx
@@ -154,7 +154,7 @@ describe('CookieConsent - CookieConsentProvider', () => {
 
     act(() => {
       result.current.saveConsent({
-        advertising: true,
+        analytics: true,
         marketing: true,
       })
     })
@@ -171,6 +171,17 @@ describe('CookieConsent - CookieConsentProvider', () => {
       maxAge: 15552000,
     })
 
+    act(() => {
+      expect(result.current.categoriesConsent).toStrictEqual({
+        analytics: true,
+        marketing: true,
+      })
+    })
+
+    act(() => {
+      expect(result.current.isSegmentAllowed).toBe(true)
+    })
+
     act(() => {
       result.current.saveConsent({
         advertising: false,