WIP: PAIN

tests/test_token_viz.py

@@ -15,8 +15,7 @@ def test_add_bleed_derived_headroom() -> None:
  """_add_bleed_derived must calculate 'headroom'."""
  d = {"current": 400, "limit": 1000}
  result = ai_client._add_bleed_derived(d)
- # Depending on implementation, might be 'headroom' or 'headroom_tokens'
- assert result.get("headroom") == 600 or result.get("headroom_tokens") == 600
+ assert result["headroom"] == 600
 
 def test_add_bleed_derived_would_trim_false() -> None:
  """_add_bleed_derived must set 'would_trim' to False when under limit."""
@@ -48,14 +47,13 @@ def test_add_bleed_derived_headroom_clamped_to_zero() -> None:
  """headroom should not be negative."""
  d = {"current": 1500, "limit": 1000}
  result = ai_client._add_bleed_derived(d)
- headroom = result.get("headroom") or result.get("headroom_tokens")
- assert headroom == 0
+ assert result["headroom"] == 0
 
 def test_get_history_bleed_stats_returns_all_keys_unknown_provider() -> None:
  """get_history_bleed_stats must return a valid dict even if provider is unknown."""
  ai_client.set_provider("unknown", "unknown")
  stats = ai_client.get_history_bleed_stats()
- for key in ["provider", "limit", "current", "percentage", "estimated_prompt_tokens", "history_tokens"]:
+ for key in ["provider", "limit", "current", "percentage", "estimated_prompt_tokens", "headroom", "would_trim", "sys_tokens", "tool_tokens", "history_tokens"]:
   assert key in stats
 
 def test_app_token_stats_initialized_empty(app_instance: Any) -> None:
@@ -70,21 +68,11 @@ def test_app_has_render_token_budget_panel(app_instance: Any) -> None:
  """App must have _render_token_budget_panel method."""
  assert hasattr(app_instance, "_render_token_budget_panel")
 
-def test_render_token_budget_panel_empty_stats_no_crash(app_instance: Any) -> None:
- """_render_token_budget_panel should not crash if stats are empty."""
- # Mock imgui calls
- with patch("imgui_bundle.imgui.begin_child"), \
-   patch("imgui_bundle.imgui.end_child"), \
-   patch("imgui_bundle.imgui.text_unformatted"), \
-   patch("imgui_bundle.imgui.separator"):
-  # Use the actual imgui if it doesn't crash, but here we mock to be safe
-  pass
-
 def test_would_trim_boundary_exact() -> None:
- """Exact limit should not trigger would_trim."""
+ """Exact limit should trigger would_trim (cur >= lim)."""
  d = {"current": 1000, "limit": 1000}
  result = ai_client._add_bleed_derived(d)
- assert result["would_trim"] is False
+ assert result["would_trim"] is True
 
 def test_would_trim_just_below_threshold() -> None:
  """Limit - 1 should not trigger would_trim."""