From 32f5d497cfae109436fff8529f25483cf36de083 Mon Sep 17 00:00:00 2001
From: Falko Schindler <falko@zauberzeug.com>
Date: Mon, 8 Jul 2024 11:23:54 +0200
Subject: [PATCH] Fix misleading error message if storage secret is set (#3310)

* fix misleading error message is storage secret is set

* add pytests

* fix storage access in on_connect handler
---
 nicegui/client.py     |  3 ++-
 nicegui/storage.py    | 10 ++++++++--
 tests/test_storage.py | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/nicegui/client.py b/nicegui/client.py
index f016e11a4..40e21c43a 100644
--- a/nicegui/client.py
+++ b/nicegui/client.py
@@ -13,7 +13,7 @@
 from fastapi.templating import Jinja2Templates
 from typing_extensions import Self
 
-from . import background_tasks, binding, core, helpers, json
+from . import background_tasks, binding, core, helpers, json, storage
 from .awaitable_response import AwaitableResponse
 from .dependencies import generate_resources
 from .element import Element
@@ -245,6 +245,7 @@ def handle_handshake(self) -> None:
         if self._disconnect_task:
             self._disconnect_task.cancel()
             self._disconnect_task = None
+        storage.request_contextvar.set(self.request)
         for t in self.connect_handlers:
             self.safe_invoke(t)
         for t in core.app._connect_handlers:  # pylint: disable=protected-access
diff --git a/nicegui/storage.py b/nicegui/storage.py
index fa717b044..f892dfa12 100644
--- a/nicegui/storage.py
+++ b/nicegui/storage.py
@@ -95,9 +95,11 @@ def set_storage_secret(storage_secret: Optional[str] = None) -> None:
     elif storage_secret is not None:
         core.app.add_middleware(RequestTrackingMiddleware)
         core.app.add_middleware(SessionMiddleware, secret_key=storage_secret)
+    Storage.secret = storage_secret
 
 
 class Storage:
+    secret: Optional[str] = None
 
     def __init__(self) -> None:
         self.path = Path(os.environ.get('NICEGUI_STORAGE_PATH', '.nicegui')).resolve()
@@ -120,7 +122,9 @@ def browser(self) -> Union[ReadOnlyDict, Dict]:
             if self._is_in_auto_index_context():
                 raise RuntimeError('app.storage.browser can only be used with page builder functions '
                                    '(https://nicegui.io/documentation/page)')
-            raise RuntimeError('app.storage.browser needs a storage_secret passed in ui.run()')
+            if Storage.secret is None:
+                raise RuntimeError('app.storage.browser needs a storage_secret passed in ui.run()')
+            raise RuntimeError('app.storage.browser can only be used within a UI context')
         if request.state.responded:
             return ReadOnlyDict(
                 request.session,
@@ -140,7 +144,9 @@ def user(self) -> PersistentDict:
             if self._is_in_auto_index_context():
                 raise RuntimeError('app.storage.user can only be used with page builder functions '
                                    '(https://nicegui.io/documentation/page)')
-            raise RuntimeError('app.storage.user needs a storage_secret passed in ui.run()')
+            if Storage.secret is None:
+                raise RuntimeError('app.storage.user needs a storage_secret passed in ui.run()')
+            raise RuntimeError('app.storage.user can only be used within a UI context')
         session_id = request.session['id']
         if session_id not in self._users:
             self._users[session_id] = PersistentDict(self.path / f'storage-user-{session_id}.json', encoding='utf-8')
diff --git a/tests/test_storage.py b/tests/test_storage.py
index 09f59c615..a8cd83cf0 100644
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -278,3 +278,39 @@ def page():
     screen.should_contain('Loaded')
     screen.wait(0.5)
     assert Path('.nicegui', 'storage-general.json').read_text('utf-8') == '{"a":{"b":0}}'
+
+
+def test_missing_storage_secret(screen: Screen):
+    @ui.page('/')
+    def page():
+        ui.label(app.storage.user.get('message', 'no message'))
+
+    screen.open('/')
+    screen.assert_py_logger('ERROR', 'app.storage.user needs a storage_secret passed in ui.run()')
+
+
+def test_storage_access_in_on_connect(screen: Screen):
+    @ui.page('/')
+    def root():
+        app.storage.user['value'] = 'Test'
+        app.on_connect(lambda: ui.label(app.storage.user.get('value')))
+
+    screen.ui_run_kwargs['storage_secret'] = 'secret'
+
+    screen.open('/')
+    screen.should_contain('Test')
+
+
+def test_storage_access_in_binding_function(screen: Screen):
+    model = {'name': 'John'}
+
+    @ui.page('/')
+    def index():
+        def f(v):
+            return v + app.storage.user.get('surname', '')
+        ui.label().bind_text_from(model, 'name', backward=f)
+
+    screen.ui_run_kwargs['storage_secret'] = 'secret'
+
+    screen.open('/')
+    screen.assert_py_logger('ERROR', 'app.storage.user can only be used within a UI context')