tests/python/unit/test_bench.py

import time

from omniff.bench.profiler import LatencyProfiler
from omniff.bench.recommend import MODELS, format_recommendation, recommend_models
from omniff.bench.suite import BenchmarkSuite


def test_benchmark_suite():
    suite = BenchmarkSuite(name="test")
    result = suite.run_benchmark(
        pipeline="text",
        fn=lambda: {"text": "hello"},
        warmup=0,
        iterations=2,
    )
    assert result.pipeline == "text"
    assert result.latency_ms >= 0


def test_benchmark_with_metric():
    suite = BenchmarkSuite()
    result = suite.run_benchmark(
        pipeline="test",
        fn=lambda: {"score": 0.95},
        metric_fn=lambda r: ("accuracy", r["score"]),
        warmup=0,
        iterations=1,
    )
    assert result.metric_name == "accuracy"
    assert result.metric_value == 0.95


def test_benchmark_summary():
    suite = BenchmarkSuite(name="test")
    suite.run_benchmark("pipe1", lambda: {}, warmup=0, iterations=1)
    summary = suite.summary()
    assert "pipe1" in summary


def test_benchmark_to_dict():
    suite = BenchmarkSuite()
    suite.run_benchmark("pipe1", lambda: {}, warmup=0, iterations=1)
    d = suite.to_dict()
    assert len(d) == 1
    assert d[0]["pipeline"] == "pipe1"


def test_profiler():
    p = LatencyProfiler()
    p.start("total")
    time.sleep(0.01)
    p.start("inner")
    time.sleep(0.01)
    p.stop()
    p.stop()
    assert len(p.entries) == 2
    assert p.entries[0].name == "inner"
    assert p.entries[1].name == "total"


def test_profiler_render():
    p = LatencyProfiler()
    p.start("op")
    p.stop()
    output = p.render()
    assert "op" in output
    assert "ms" in output


def test_profiler_empty():
    p = LatencyProfiler()
    assert p.render() == "(no profile data)"


def test_recommend_small_budget():
    models = recommend_models(5.0, ["text"])
    assert len(models) >= 1
    total = sum(m.vram_gb for m in models)
    assert total <= 5.0


def test_recommend_large_budget():
    models = recommend_models(40.0, ["text", "image", "audio"])
    assert len(models) >= 2


def test_recommend_specific_pipelines():
    models = recommend_models(10.0, ["audio"])
    for m in models:
        assert "audio" in m.pipelines


def test_format_recommendation():
    models = recommend_models(22.0)
    output = format_recommendation(models, 22.0)
    assert "22" in output
    assert "VRAM" in output


def test_models_catalog():
    assert len(MODELS) >= 5
    for m in MODELS:
        assert m.vram_gb > 0
        assert len(m.pipelines) > 0