mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 14:07:55 +00:00
79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
"""Tests for the throughput benchmark."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from openjarvis.bench.throughput import ThroughputBenchmark
|
|
from openjarvis.core.registry import BenchmarkRegistry
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _register_throughput():
|
|
"""Re-register throughput benchmark after registry clear."""
|
|
from openjarvis.bench.throughput import ensure_registered
|
|
|
|
ensure_registered()
|
|
|
|
|
|
def _make_engine(completion_tokens=10):
|
|
engine = MagicMock()
|
|
engine.engine_id = "mock"
|
|
engine.generate.return_value = {
|
|
"content": "Hello world",
|
|
"usage": {
|
|
"prompt_tokens": 5,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": 5 + completion_tokens,
|
|
},
|
|
}
|
|
return engine
|
|
|
|
|
|
class TestThroughputBenchmark:
|
|
def test_registration(self):
|
|
assert BenchmarkRegistry.contains("throughput")
|
|
assert BenchmarkRegistry.get("throughput") is ThroughputBenchmark
|
|
|
|
def test_run_with_mock(self):
|
|
engine = _make_engine()
|
|
b = ThroughputBenchmark()
|
|
result = b.run(engine, "test-model", num_samples=3)
|
|
assert result.benchmark_name == "throughput"
|
|
assert result.model == "test-model"
|
|
assert result.engine == "mock"
|
|
assert result.samples == 3
|
|
|
|
def test_metrics_keys(self):
|
|
engine = _make_engine()
|
|
b = ThroughputBenchmark()
|
|
result = b.run(engine, "test-model", num_samples=3)
|
|
assert "mean_tokens_per_second" in result.metrics
|
|
assert "total_tokens" in result.metrics
|
|
assert "total_time_seconds" in result.metrics
|
|
|
|
def test_tokens_per_second_calc(self):
|
|
engine = _make_engine(completion_tokens=10)
|
|
b = ThroughputBenchmark()
|
|
result = b.run(engine, "test-model", num_samples=5)
|
|
# 5 samples * 10 tokens each = 50 total tokens
|
|
assert result.metrics["total_tokens"] == 50.0
|
|
assert result.metrics["mean_tokens_per_second"] > 0
|
|
|
|
def test_sample_count(self):
|
|
engine = _make_engine()
|
|
b = ThroughputBenchmark()
|
|
b.run(engine, "test-model", num_samples=7)
|
|
assert engine.generate.call_count == 7
|
|
|
|
def test_zero_latency_handling(self):
|
|
"""All errors should result in 0 tokens_per_second."""
|
|
engine = _make_engine()
|
|
engine.generate.side_effect = RuntimeError("fail")
|
|
b = ThroughputBenchmark()
|
|
result = b.run(engine, "test-model", num_samples=3)
|
|
assert result.errors == 3
|
|
assert result.metrics.get("mean_tokens_per_second", 0.0) == 0.0
|