From 8503a273796f5cb2a0bd97c849cf3875eeca5baa Mon Sep 17 00:00:00 2001
From: pyr0ball <pyroballpcs@gmail.com>
Date: Sun, 22 Mar 2026 16:27:20 -0700
Subject: [PATCH] fix(demo): remove reload antipattern, fix label consistency
 in demo_toolbar tests

---
 app/components/demo_toolbar.py |  6 +-----
 tests/test_demo_toolbar.py     | 30 ++++++++++++------------------
 2 files changed, 13 insertions(+), 23 deletions(-)

diff --git a/app/components/demo_toolbar.py b/app/components/demo_toolbar.py
index 829c25c..2c30c56 100644
--- a/app/components/demo_toolbar.py
+++ b/app/components/demo_toolbar.py
@@ -50,11 +50,7 @@ def render_demo_toolbar() -> None:
     """
     current = get_simulated_tier()
 
-    labels = {
-        "free": "Free",
-        "paid": "Paid ✓" if current == "paid" else "Paid",
-        "premium": "Premium ✓" if current == "premium" else "Premium",
-    }
+    labels = {t: t.capitalize() + (" ✓" if t == current else "") for t in _VALID_TIERS}
 
     with st.container():
         cols = st.columns([3, 1, 1, 1, 2])
diff --git a/tests/test_demo_toolbar.py b/tests/test_demo_toolbar.py
index e2289ba..c7cb155 100644
--- a/tests/test_demo_toolbar.py
+++ b/tests/test_demo_toolbar.py
@@ -1,12 +1,15 @@
 """Tests for app/components/demo_toolbar.py."""
-import sys, os
+import sys
 from pathlib import Path
 import pytest
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
-# Ensure DEMO_MODE is set so the module initialises correctly
-os.environ["DEMO_MODE"] = "true"
+from app.components.demo_toolbar import (
+    get_simulated_tier,
+    set_simulated_tier,
+    render_demo_toolbar,
+)
 
 
 def test_set_simulated_tier_updates_session_state(monkeypatch):
@@ -17,12 +20,7 @@ def test_set_simulated_tier_updates_session_state(monkeypatch):
     monkeypatch.setattr("streamlit.session_state", session, raising=False)
     monkeypatch.setattr("streamlit.rerun", lambda: None)
 
-    from unittest.mock import patch
-    with patch('app.components.demo_toolbar._DEMO_MODE', True):
-        from importlib import reload
-        import app.components.demo_toolbar as m
-        reload(m)
-        m.set_simulated_tier("premium")
+    set_simulated_tier("premium")
 
     assert session.get("simulated_tier") == "premium"
     assert any("prgn_demo_tier=premium" in h for h in injected)
@@ -35,12 +33,7 @@ def test_set_simulated_tier_invalid_ignored(monkeypatch):
     monkeypatch.setattr("streamlit.session_state", session, raising=False)
     monkeypatch.setattr("streamlit.rerun", lambda: None)
 
-    from unittest.mock import patch
-    with patch('app.components.demo_toolbar._DEMO_MODE', True):
-        from importlib import reload
-        import app.components.demo_toolbar as m
-        reload(m)
-        m.set_simulated_tier("ultramax")
+    set_simulated_tier("ultramax")
 
     assert "simulated_tier" not in session
 
@@ -50,7 +43,6 @@ def test_get_simulated_tier_defaults_to_paid(monkeypatch):
     monkeypatch.setattr("streamlit.session_state", {}, raising=False)
     monkeypatch.setattr("streamlit.query_params", {}, raising=False)
 
-    from app.components.demo_toolbar import get_simulated_tier
     assert get_simulated_tier() == "paid"
 
 
@@ -59,7 +51,6 @@ def test_get_simulated_tier_reads_session(monkeypatch):
     monkeypatch.setattr("streamlit.session_state", {"simulated_tier": "free"}, raising=False)
     monkeypatch.setattr("streamlit.query_params", {}, raising=False)
 
-    from app.components.demo_toolbar import get_simulated_tier
     assert get_simulated_tier() == "free"
 
 
@@ -79,10 +70,13 @@ def test_render_demo_toolbar_renders_pills(monkeypatch):
     monkeypatch.setattr("streamlit.button", mock_button)
     monkeypatch.setattr("streamlit.divider", lambda: None)
 
-    from app.components.demo_toolbar import render_demo_toolbar
     render_demo_toolbar()
 
     # Verify buttons were rendered for all tiers
     button_calls = [c for c in calls if c[0] == "button"]
     assert len(button_calls) == 3
     assert any("Paid ✓" in c[1] for c in button_calls)  # current tier marked
+
+    primary_calls = [c for c in button_calls if c[3] == "primary"]
+    assert len(primary_calls) == 1
+    assert "Paid" in primary_calls[0][1]