161 lines
4.8 KiB
Python
161 lines
4.8 KiB
Python
|
|
from datetime import datetime, timedelta
|
||
|
|
from types import SimpleNamespace
|
||
|
|
|
||
|
|
from cli import HermesCLI
|
||
|
|
|
||
|
|
|
||
|
|
def _make_cli(model: str = "anthropic/claude-sonnet-4-20250514"):
|
||
|
|
cli_obj = HermesCLI.__new__(HermesCLI)
|
||
|
|
cli_obj.model = model
|
||
|
|
cli_obj.session_start = datetime.now() - timedelta(minutes=14, seconds=32)
|
||
|
|
cli_obj.conversation_history = [{"role": "user", "content": "hi"}]
|
||
|
|
cli_obj.agent = None
|
||
|
|
return cli_obj
|
||
|
|
|
||
|
|
|
||
|
|
def _attach_agent(
|
||
|
|
cli_obj,
|
||
|
|
*,
|
||
|
|
prompt_tokens: int,
|
||
|
|
completion_tokens: int,
|
||
|
|
total_tokens: int,
|
||
|
|
api_calls: int,
|
||
|
|
context_tokens: int,
|
||
|
|
context_length: int,
|
||
|
|
compressions: int = 0,
|
||
|
|
):
|
||
|
|
cli_obj.agent = SimpleNamespace(
|
||
|
|
model=cli_obj.model,
|
||
|
|
session_prompt_tokens=prompt_tokens,
|
||
|
|
session_completion_tokens=completion_tokens,
|
||
|
|
session_total_tokens=total_tokens,
|
||
|
|
session_api_calls=api_calls,
|
||
|
|
context_compressor=SimpleNamespace(
|
||
|
|
last_prompt_tokens=context_tokens,
|
||
|
|
context_length=context_length,
|
||
|
|
compression_count=compressions,
|
||
|
|
),
|
||
|
|
)
|
||
|
|
return cli_obj
|
||
|
|
|
||
|
|
|
||
|
|
class TestCLIStatusBar:
|
||
|
|
def test_context_style_thresholds(self):
|
||
|
|
cli_obj = _make_cli()
|
||
|
|
|
||
|
|
assert cli_obj._status_bar_context_style(None) == "class:status-bar-dim"
|
||
|
|
assert cli_obj._status_bar_context_style(10) == "class:status-bar-good"
|
||
|
|
assert cli_obj._status_bar_context_style(50) == "class:status-bar-warn"
|
||
|
|
assert cli_obj._status_bar_context_style(81) == "class:status-bar-bad"
|
||
|
|
assert cli_obj._status_bar_context_style(95) == "class:status-bar-critical"
|
||
|
|
|
||
|
|
def test_build_status_bar_text_for_wide_terminal(self):
|
||
|
|
cli_obj = _attach_agent(
|
||
|
|
_make_cli(),
|
||
|
|
prompt_tokens=10_230,
|
||
|
|
completion_tokens=2_220,
|
||
|
|
total_tokens=12_450,
|
||
|
|
api_calls=7,
|
||
|
|
context_tokens=12_450,
|
||
|
|
context_length=200_000,
|
||
|
|
)
|
||
|
|
|
||
|
|
text = cli_obj._build_status_bar_text(width=120)
|
||
|
|
|
||
|
|
assert "claude-sonnet-4-20250514" in text
|
||
|
|
assert "12.4K/200K" in text
|
||
|
|
assert "6%" in text
|
||
|
|
assert "$0.06" in text
|
||
|
|
assert "15m" in text
|
||
|
|
|
||
|
|
def test_build_status_bar_text_collapses_for_narrow_terminal(self):
|
||
|
|
cli_obj = _attach_agent(
|
||
|
|
_make_cli(),
|
||
|
|
prompt_tokens=10_230,
|
||
|
|
completion_tokens=2_220,
|
||
|
|
total_tokens=12_450,
|
||
|
|
api_calls=7,
|
||
|
|
context_tokens=12_450,
|
||
|
|
context_length=200_000,
|
||
|
|
)
|
||
|
|
|
||
|
|
text = cli_obj._build_status_bar_text(width=60)
|
||
|
|
|
||
|
|
assert "⚕" in text
|
||
|
|
assert "$0.06" in text
|
||
|
|
assert "15m" in text
|
||
|
|
assert "200K" not in text
|
||
|
|
|
||
|
|
def test_build_status_bar_text_handles_missing_agent(self):
|
||
|
|
cli_obj = _make_cli()
|
||
|
|
|
||
|
|
text = cli_obj._build_status_bar_text(width=100)
|
||
|
|
|
||
|
|
assert "⚕" in text
|
||
|
|
assert "claude-sonnet-4-20250514" in text
|
||
|
|
|
||
|
|
|
||
|
|
class TestCLIUsageReport:
|
||
|
|
def test_show_usage_includes_estimated_cost(self, capsys):
|
||
|
|
cli_obj = _attach_agent(
|
||
|
|
_make_cli(),
|
||
|
|
prompt_tokens=10_230,
|
||
|
|
completion_tokens=2_220,
|
||
|
|
total_tokens=12_450,
|
||
|
|
api_calls=7,
|
||
|
|
context_tokens=12_450,
|
||
|
|
context_length=200_000,
|
||
|
|
compressions=1,
|
||
|
|
)
|
||
|
|
cli_obj.verbose = False
|
||
|
|
|
||
|
|
cli_obj._show_usage()
|
||
|
|
output = capsys.readouterr().out
|
||
|
|
|
||
|
|
assert "Model:" in output
|
||
|
|
assert "Input cost:" in output
|
||
|
|
assert "Output cost:" in output
|
||
|
|
assert "Total cost:" in output
|
||
|
|
assert "$" in output
|
||
|
|
assert "0.064" in output
|
||
|
|
assert "Session duration:" in output
|
||
|
|
assert "Compressions:" in output
|
||
|
|
|
||
|
|
def test_show_usage_marks_unknown_pricing(self, capsys):
|
||
|
|
cli_obj = _attach_agent(
|
||
|
|
_make_cli(model="local/my-custom-model"),
|
||
|
|
prompt_tokens=1_000,
|
||
|
|
completion_tokens=500,
|
||
|
|
total_tokens=1_500,
|
||
|
|
api_calls=1,
|
||
|
|
context_tokens=1_000,
|
||
|
|
context_length=32_000,
|
||
|
|
)
|
||
|
|
cli_obj.verbose = False
|
||
|
|
|
||
|
|
cli_obj._show_usage()
|
||
|
|
output = capsys.readouterr().out
|
||
|
|
|
||
|
|
assert "Total cost:" in output
|
||
|
|
assert "n/a" in output
|
||
|
|
assert "Pricing unknown for local/my-custom-model" in output
|
||
|
|
|
||
|
|
def test_zero_priced_provider_models_stay_unknown(self, capsys):
|
||
|
|
cli_obj = _attach_agent(
|
||
|
|
_make_cli(model="glm-5"),
|
||
|
|
prompt_tokens=1_000,
|
||
|
|
completion_tokens=500,
|
||
|
|
total_tokens=1_500,
|
||
|
|
api_calls=1,
|
||
|
|
context_tokens=1_000,
|
||
|
|
context_length=32_000,
|
||
|
|
)
|
||
|
|
cli_obj.verbose = False
|
||
|
|
|
||
|
|
cli_obj._show_usage()
|
||
|
|
output = capsys.readouterr().out
|
||
|
|
|
||
|
|
assert "Total cost:" in output
|
||
|
|
assert "n/a" in output
|
||
|
|
assert "Pricing unknown for glm-5" in output
|