From 8eed0f4e7362b0543d534a8093be0fc0d8f64cce Mon Sep 17 00:00:00 2001
From: Joseph Silcock <55981653+josephsilcock@users.noreply.github.com>
Date: Tue, 29 Oct 2024 05:19:12 +0000
Subject: [PATCH] fix cookie not loading on initial load (#14)

---
 .../addon/src/components/PanelContent.tsx     | 42 ++++++++-----------
 packages/addon/src/constants.ts               |  2 +-
 2 files changed, 18 insertions(+), 26 deletions(-)

diff --git a/packages/addon/src/components/PanelContent.tsx b/packages/addon/src/components/PanelContent.tsx
index dd74d57..f8fbc56 100644
--- a/packages/addon/src/components/PanelContent.tsx
+++ b/packages/addon/src/components/PanelContent.tsx
@@ -1,49 +1,41 @@
 import React, { useEffect, useState } from 'react';
-import {
-  useGlobals,
-  useParameter,
-  useStorybookState,
-} from '@storybook/manager-api';
+import { useGlobals, useParameter } from '@storybook/manager-api';
 import Icons from '@storybook/icons';
 import { ObjectControl } from '@storybook/blocks';
 
-import { PARAM_ENCONDING_KEY, PARAM_KEY } from '../constants';
+import { PARAM_ENCODING_KEY, PARAM_KEY } from '../constants';
 import { Cookie } from '../types';
 import { clearCookies, setCookies } from '../utils';
 
 export const PanelContent: React.FC = () => {
-  const { path } = useStorybookState();
-  const defaultCookie = useParameter<Cookie>(PARAM_KEY, {}); // TODO: Doesn't work on initial load
-  const encoding = useParameter<boolean>(PARAM_ENCONDING_KEY, false);
+  const defaultCookie = useParameter<Cookie>(PARAM_KEY, null);
+  const encoding = useParameter<boolean>(PARAM_ENCODING_KEY, false);
 
-  const [story, setStory] = useState<string>();
   const [value, setValue] = useState<Cookie>();
   const [globals, updateGlobals] = useGlobals();
 
-  const initCookie = () => {
-    setStory(path);
-    setValue(defaultCookie);
+  const updateCookieValue = (newValue: Cookie | null) => {
+    setValue(newValue ?? {});
     clearCookies();
-    setCookies(defaultCookie, encoding);
+    if (newValue) {
+      setCookies(newValue, encoding);
+    }
+    updateGlobals({ ...globals });
   };
 
   useEffect(() => {
-    if (story !== path) {
-      initCookie();
+    if (!defaultCookie && !value) {
+      return;
     }
-  }, [path, story]);
+    updateCookieValue(defaultCookie);
+  }, [defaultCookie, encoding]);
 
   const handleChange = (newValue: Cookie) => {
-    clearCookies();
-    setCookies(newValue, encoding);
-    setValue(newValue);
-    updateGlobals({ ...globals });
+    updateCookieValue(newValue);
   };
 
   const handleClear = () => {
-    clearCookies();
-    setValue({});
-    updateGlobals({ ...globals });
+    updateCookieValue(null);
   };
 
   return (
@@ -55,7 +47,7 @@ export const PanelContent: React.FC = () => {
       <ObjectControl
         name="cookie"
         onChange={handleChange}
-        value={value || defaultCookie}
+        value={value}
         theme=""
       />
     </div>
diff --git a/packages/addon/src/constants.ts b/packages/addon/src/constants.ts
index 73c17ce..ffedf8b 100644
--- a/packages/addon/src/constants.ts
+++ b/packages/addon/src/constants.ts
@@ -2,5 +2,5 @@ export const ADDON_ID = 'storybook-addon-cookie';
 export const PANEL_ID = `${ADDON_ID}/panel`;
 export const PARAM_KEY = `cookie`;
 export const ADDON_TITLE = 'Cookie';
-export const PARAM_ENCONDING_KEY = 'cookieEncoding';
+export const PARAM_ENCODING_KEY = 'cookieEncoding';
 export const PARAM_PRESERVE_KEY = 'cookiePreserve';