homelab-codex-ws/jobs/mail-body-ingest/tests/test_benchmark.py

"""Unit tests for the batch-size benchmark — no real HTTP, no real Ollama, no DB."""
from __future__ import annotations

import aiohttp
import pytest

from mail_body_ingest.benchmark import (
    FULL_CORPUS_CHUNKS,
    BenchRow,
    bench_size,
    format_table,
    split_batches,
)


class TestSplitBatches:
    def test_exact_division(self):
        assert split_batches([1, 2, 3, 4], 2) == [[1, 2], [3, 4]]

    def test_trailing_short_batch_is_kept(self):
        assert split_batches([1, 2, 3, 4, 5], 2) == [[1, 2], [3, 4], [5]]

    def test_size_one_is_one_item_per_batch(self):
        assert split_batches([1, 2, 3], 1) == [[1], [2], [3]]

    def test_size_larger_than_input_yields_a_single_batch(self):
        assert split_batches([1, 2], 64) == [[1, 2]]

    def test_empty_input_yields_no_batches(self):
        assert split_batches([], 8) == []

    def test_covers_every_item_exactly_once(self):
        items = list(range(100))
        flattened = [i for batch in split_batches(items, 7) for i in batch]
        assert flattened == items

    def test_zero_or_negative_size_raises(self):
        with pytest.raises(ValueError):
            split_batches([1, 2], 0)


class TestBenchRow:
    def test_ms_per_chunk(self):
        row = BenchRow(batch_size=64, chunks=100, requests=2, seconds=2.0)
        assert row.ms_per_chunk == 20.0

    def test_chunks_per_second(self):
        row = BenchRow(batch_size=64, chunks=100, requests=2, seconds=2.0)
        assert row.chunks_per_s == 50.0

    def test_projected_hours_scales_the_measured_rate_to_the_full_corpus(self):
        row = BenchRow(batch_size=64, chunks=100, requests=2, seconds=1.0)  # 10 ms/chunk
        assert row.projected_hours == pytest.approx(FULL_CORPUS_CHUNKS * 0.010 / 3600)

    def test_zero_chunks_does_not_divide_by_zero(self):
        row = BenchRow(batch_size=1, chunks=0, requests=0, seconds=0.0)
        assert row.ms_per_chunk == 0.0
        assert row.chunks_per_s == 0.0


class TestFormatTable:
    def _rows(self):
        return [
            BenchRow(batch_size=1, chunks=100, requests=100, seconds=20.0),   # 200 ms/chunk
            BenchRow(batch_size=64, chunks=100, requests=2, seconds=1.0),     # 10 ms/chunk
        ]

    def test_one_line_per_row_plus_header(self):
        out = format_table(self._rows())
        assert "batch" in out and "ms/chunk" in out
        assert "200.00" in out and "10.00" in out

    def test_names_the_best_size_and_the_speedup_over_batch_1(self):
        out = format_table(self._rows())
        assert "best: batch=64" in out
        assert "20.0x faster than batch=1" in out

    def test_no_speedup_note_when_batch_1_wins(self):
        rows = [BenchRow(batch_size=1, chunks=100, requests=100, seconds=1.0)]
        out = format_table(rows)
        assert "best: batch=1" in out
        assert "faster than" not in out

    def test_empty_rows_still_renders_a_header(self):
        assert "batch" in format_table([])

    def test_a_size_that_failed_entirely_is_not_reported_as_best(self):
        """A batch size the backend chokes on measures 0 chunks; ranking it best would read as
        'infinitely fast' instead of 'did not work'."""
        rows = [
            BenchRow(batch_size=8, chunks=100, requests=13, seconds=2.0),
            BenchRow(batch_size=512, chunks=0, requests=0, seconds=0.0, failures=1),
        ]
        assert "best: batch=8" in format_table(rows)


class _FakeResponse:
    def __init__(self, payload, fail=False):
        self._payload = payload
        self._fail = fail

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False

    async def json(self):
        return self._payload

    def raise_for_status(self):
        if self._fail:
            raise aiohttp.ClientConnectionError("simulated failure")


class _FakeSession:
    def __init__(self, dim=1024, fail_sizes=()):
        self._dim = dim
        self._fail_sizes = set(fail_sizes)
        self.posts: list[int] = []

    def post(self, url, json, timeout=None):
        texts = json["input"]
        self.posts.append(len(texts))
        if len(texts) in self._fail_sizes:
            return _FakeResponse({}, fail=True)
        return _FakeResponse({"embeddings": [[0.01] * self._dim for _ in texts]})


class TestBenchSize:
    async def test_measures_every_chunk_once(self):
        session = _FakeSession()
        chunks = [f"c{i}" for i in range(10)]
        row = await bench_size(session, "http://fake", "bge-m3", chunks, 4, 60.0)
        assert row.chunks == 10
        assert row.requests == 3          # 4 + 4 + 2
        assert session.posts == [4, 4, 2]
        assert row.failures == 0

    async def test_failures_are_counted_not_raised(self):
        """A size the backend can't serve is a result, not a reason to abandon the sweep."""
        session = _FakeSession(fail_sizes={8})
        chunks = [f"c{i}" for i in range(16)]
        row = await bench_size(session, "http://fake", "bge-m3", chunks, 8, 60.0)
        assert row.failures == 2
        assert row.chunks == 0
        assert row.ms_per_chunk == 0.0

    async def test_partial_failure_still_measures_the_successful_batches(self):
        session = _FakeSession(fail_sizes={2})  # only the short trailing batch fails
        chunks = [f"c{i}" for i in range(10)]
        row = await bench_size(session, "http://fake", "bge-m3", chunks, 4, 60.0)
        assert row.chunks == 8
        assert row.requests == 2
        assert row.failures == 1