From e3a1311f32bb4cd220bfe8735586d7ddf600d203 Mon Sep 17 00:00:00 2001
From: Aarthy Adityan <aarthy@datakitchen.io>
Date: Thu, 10 Oct 2024 12:28:06 -0400
Subject: [PATCH] fix(auth): handle cookies retrieved inconsistently

---
 testgen/ui/navigation/router.py | 15 +++++++++++----
 testgen/ui/session.py           |  2 +-
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/testgen/ui/navigation/router.py b/testgen/ui/navigation/router.py
index d010ee9..011ebb8 100644
--- a/testgen/ui/navigation/router.py
+++ b/testgen/ui/navigation/router.py
@@ -10,6 +10,7 @@
 from testgen.utils.singleton import Singleton
 
 LOG = logging.getLogger("testgen")
+COOKIES_READY_RERUNS = 2
 
 
 class Router(Singleton):
@@ -32,12 +33,16 @@ def run(self, hide_sidebar=False) -> None:
         session.current_page_args = st.query_params
 
         # This hack is needed because the auth cookie is not retrieved on the first run
-        # We have to store the page and wait for the second run
-
+        # We have to store the page and wait for the second or third run
         if not session.cookies_ready:
-            session.cookies_ready = True
+            session.cookies_ready = 1
             session.page_pending_cookies = current_page
-        else:
+
+        # Sometimes the cookie is ready on the second rerun and other times only on the third -_-
+        # so we have to make sure the page renders correctly in both cases
+        # and also handle the login page!
+        elif session.cookies_ready == COOKIES_READY_RERUNS or session.authentication_status or (session.page_pending_cookies and not session.page_pending_cookies.url_path):
+            session.cookies_ready = COOKIES_READY_RERUNS
             current_page = session.page_pending_cookies or current_page
             session.page_pending_cookies = None
 
@@ -48,6 +53,8 @@ def run(self, hide_sidebar=False) -> None:
 
             session.current_page = current_page.url_path
             current_page.run()
+        else:
+            session.cookies_ready += 1
 
 
     def navigate(self, /, to: str, with_args: dict = {}) -> None:  # noqa: B006
diff --git a/testgen/ui/session.py b/testgen/ui/session.py
index 5c7459a..0802132 100644
--- a/testgen/ui/session.py
+++ b/testgen/ui/session.py
@@ -7,7 +7,7 @@
 
 
 class TestgenSession(Singleton):
-    cookies_ready: bool
+    cookies_ready: int
     logging_in: bool
     logging_out: bool
     page_pending_cookies: st.Page