Function _make_batch_provider

tests/extraction/test_runner.py:101–118

Build a mock batch provider.

(
    responses: dict[str, str] | None = None,
    error: Exception | None = None,
)

def _make_batch_provider(
    responses: dict[str, str] | None = None,
    error: Exception | None = None,
) -> MagicMock:
    """Build a mock batch provider."""
    bp = MagicMock()
    bp.make_poll_client.return_value = MagicMock()

    bp.submit_batch.return_value = "test-job-id"
    bp.poll_batch.return_value = MagicMock()

    if error:
        bp.collect_results.side_effect = error
    else:
        usage = TokenUsage(input_tokens=500, output_tokens=200)
        bp.collect_results.return_value = BatchResults(responses or {}, usage)

    return bp


class TestRunBatchStatsContract(unittest.TestCase):