111 lines
4.5 KiB
Python
111 lines
4.5 KiB
Python
"""Tests for context & token visualization (Track: context_token_viz_20260301)."""
|
|
|
|
from src import ai_client
|
|
from typing import Any
|
|
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
def test_add_bleed_derived_aliases() -> None:
|
|
"""_add_bleed_derived must inject 'estimated_prompt_tokens' alias."""
|
|
d = {"current": 100, "limit": 1000}
|
|
result = ai_client._add_bleed_derived(d)
|
|
assert result["estimated_prompt_tokens"] == 100
|
|
|
|
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
|
|
|
|
def test_add_bleed_derived_would_trim_false() -> None:
|
|
"""_add_bleed_derived must set 'would_trim' to False when under limit."""
|
|
d = {"current": 100, "limit": 1000}
|
|
result = ai_client._add_bleed_derived(d)
|
|
assert result["would_trim"] is False
|
|
|
|
def test_add_bleed_derived_would_trim_true() -> None:
|
|
"""_add_bleed_derived must set 'would_trim' to True when over limit."""
|
|
d = {"current": 1100, "limit": 1000}
|
|
result = ai_client._add_bleed_derived(d)
|
|
assert result["would_trim"] is True
|
|
|
|
def test_add_bleed_derived_breakdown() -> None:
|
|
"""_add_bleed_derived must calculate breakdown of current usage."""
|
|
d = {"current": 500, "limit": 1000}
|
|
result = ai_client._add_bleed_derived(d, sys_tok=100, tool_tok=50)
|
|
assert result["sys_tokens"] == 100
|
|
assert result["tool_tokens"] == 50
|
|
assert result["history_tokens"] == 350
|
|
|
|
def test_add_bleed_derived_history_clamped_to_zero() -> None:
|
|
"""history_tokens should not be negative."""
|
|
d = {"current": 50, "limit": 1000}
|
|
result = ai_client._add_bleed_derived(d, sys_tok=100, tool_tok=50)
|
|
assert result["history_tokens"] == 0
|
|
|
|
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
|
|
|
|
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"]:
|
|
assert key in stats
|
|
|
|
def test_app_token_stats_initialized_empty(app_instance: Any) -> None:
|
|
"""App._token_stats should start empty."""
|
|
assert app_instance.controller._token_stats == {}
|
|
|
|
def test_app_last_stable_md_initialized_empty(app_instance: Any) -> None:
|
|
"""App._last_stable_md should start empty."""
|
|
assert app_instance.controller._last_stable_md == ''
|
|
|
|
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."""
|
|
d = {"current": 1000, "limit": 1000}
|
|
result = ai_client._add_bleed_derived(d)
|
|
assert result["would_trim"] is False
|
|
|
|
def test_would_trim_just_below_threshold() -> None:
|
|
"""Limit - 1 should not trigger would_trim."""
|
|
d = {"current": 999, "limit": 1000}
|
|
result = ai_client._add_bleed_derived(d)
|
|
assert result["would_trim"] is False
|
|
|
|
def test_would_trim_just_above_threshold() -> None:
|
|
"""Limit + 1 should trigger would_trim."""
|
|
d = {"current": 1001, "limit": 1000}
|
|
result = ai_client._add_bleed_derived(d)
|
|
assert result["would_trim"] is True
|
|
|
|
def test_gemini_cache_fields_accessible() -> None:
|
|
"""_gemini_cache and related fields must be accessible for stats rendering."""
|
|
assert hasattr(ai_client, "_gemini_cache")
|
|
assert hasattr(ai_client, "_gemini_cache_created_at")
|
|
assert hasattr(ai_client, "_GEMINI_CACHE_TTL")
|
|
|
|
def test_anthropic_history_lock_accessible() -> None:
|
|
"""_anthropic_history_lock must be accessible for cache hint rendering."""
|
|
assert hasattr(ai_client, "_anthropic_history_lock")
|
|
assert hasattr(ai_client, "_anthropic_history")
|