wehub-resource-sync
33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
"""Tests for output formatting helpers."""
|
|
|
|
from whichllm.engine.types import CompatibilityResult
|
|
from whichllm.models.types import ModelInfo
|
|
from whichllm.output.formatting import _format_speed
|
|
|
|
|
|
def _result(speed: float, confidence: str) -> CompatibilityResult:
|
|
return CompatibilityResult(
|
|
model=ModelInfo(
|
|
id="org/model",
|
|
family_id="org/model",
|
|
name="model",
|
|
parameter_count=1_000_000_000,
|
|
),
|
|
gguf_variant=None,
|
|
can_run=True,
|
|
vram_required_bytes=0,
|
|
vram_available_bytes=0,
|
|
estimated_tok_per_sec=speed,
|
|
speed_confidence=confidence,
|
|
)
|
|
|
|
|
|
def test_format_speed_colors_by_runtime_speed_not_confidence():
|
|
assert _format_speed(_result(2.5, "medium")) == "[red]2.5 tok/s ~[/red]"
|
|
assert _format_speed(_result(6.0, "low")) == "[yellow]6.0 tok/s ?[/yellow]"
|
|
assert _format_speed(_result(12.0, "medium")) == "[green]12.0 tok/s ~[/green]"
|
|
assert (
|
|
_format_speed(_result(30.0, "low"))
|
|
== "[bright_green]30.0 tok/s ?[/bright_green]"
|
|
)
|