Testing Partial Failures in a Batch Inference Job

A batch of a hundred documents goes to a model. Row 37 comes back with a 400, row 61 times out, and the job throws. Ninety-eight good results are discarded because two were bad. The test that would have caught this does not assert on the model at all.

The envelope is the unit under test

The bug is almost always one line, and it is usually Promise.all or a bare list comprehension around an await. Both are all-or-nothing by design: the first rejection becomes the result of the whole thing, and the work already done is unreachable because nothing kept a reference to it.

So the thing to test is not the provider call. It is the shape your batch function returns. A batch that can partially fail has to return one entry per input row, each of which is either a success or a located failure — a discriminated union, one per row, with the input index carried through. Once that shape exists, the test writes itself, because every interesting property is a property of an array you can inspect without a network.

// batch.ts
export type RowResult<T> =
  | { index: number; ok: true; value: T }
  | { index: number; ok: false; error: string; retryable: boolean };

export async function runBatch<I, T>(
  rows: I[],
  call: (row: I) => Promise<T>,
  concurrency = 8,
): Promise<RowResult<T>[]> {
  const out: RowResult<T>[] = new Array(rows.length);
  let next = 0;
  const worker = async () => {
    while (true) {
      const i = next++;
      if (i >= rows.length) return;
      try {
        out[i] = { index: i, ok: true, value: await call(rows[i]) };
      } catch (e) {
        out[i] = {
          index: i,
          ok: false,
          error: e instanceof Error ? e.message : String(e),
          retryable: isRetryable(e),
        };
      }
    }
  };
  await Promise.all(Array.from({ length: concurrency }, worker));
  return out;
}

Note where Promise.all survives: it waits on the workers, not on the rows. A worker never rejects, because every call is wrapped. That is the whole mechanism, and it is why the assertions below can be strict.

A fake provider that fails one row

You do not need a mocked HTTP layer for this. The batch function takes the per-row call as an argument precisely so the test can supply a function that fails deterministically on the rows you choose. Keep the fake dumb: a lookup from input to outcome, no timing, no randomness.

// batch.test.ts
import { describe, it, expect, vi } from "vitest";
import { runBatch } from "./batch";

const rows = ["a", "b", "c", "d", "e"];

const flaky = vi.fn(async (row: string) => {
  if (row === "c") throw new Error("400 invalid_request_error: image too large");
  if (row === "d") throw Object.assign(new Error("429 rate_limit_exceeded"), { status: 429 });
  return row.toUpperCase();
});

describe("runBatch", () => {
  it("returns one entry per input row even when two fail", async () => {
    const results = await runBatch(rows, flaky, 2);
    expect(results).toHaveLength(rows.length);
    expect(results.map((r) => r.index)).toEqual([0, 1, 2, 3, 4]);
  });
});

Two rows fail for different reasons on purpose. A single failing row proves less than you think: it cannot distinguish a function that reports failures from one that stops at the first and pads the rest.

Keep the fake synchronous in spirit even though it is declared async. Randomised failures make the suite flaky and, worse, make a failure hard to reproduce from the output; if you want to model a rate of failure rather than specific rows, drive it from a seeded sequence and print the seed. The one thing worth making asynchronous on purpose is timing, and only in the ordering test below, where the whole point is that completion order and input order come apart.

What to assert

Five assertions, in rough order of how often each one catches something real.

• The call does not reject. await expect(runBatch(...)).resolves.toBeDefined(). If a partial failure can still throw, nothing else matters.
• Length equals input length. A result array shorter than the input is the silent version of this bug: no exception, no error log, just rows missing from the output file.
• Every successful row has the value it would have had alone. Failures must not perturb neighbours. Assert the actual values, not just ok: true.
• Each failure names its row and classifies its error. The index is what makes a partial failure actionable, and the retryable flag is what the next section depends on. Assert the 429 is retryable and the 400 is not.
• The provider was called once per row. expect(flaky).toHaveBeenCalledTimes(5) catches a short-circuit that a length assertion misses, because a short-circuiting implementation can still fill the array with placeholder failures.

What is not on the list: anything about the text the model produced. The successful rows are uppercase strings from a fake. This test is about control flow, and mixing an output-quality assertion into it gives you a test that goes red for two unrelated reasons.

Ordering, and the off-by-one it hides

Under concurrency, completion order is not input order. A worker pool finishes row 4 before row 2 routinely. If your implementation collects results by pushing onto an array as they complete, the output is in completion order — and then row 37 of the output is not row 37 of the input. Every downstream join by position is now wrong, and nothing throws.

This is why the fake above should not be uniform. Make one row slow and assert positional identity explicitly:

it("keeps results in input order regardless of completion order", async () => {
  const call = async (row: string) => {
    if (row === "a") await new Promise((r) => setTimeout(r, 20));
    return row.toUpperCase();
  };
  const results = await runBatch(rows, call, 5);
  expect(results.filter((r) => r.ok).map((r) => (r as { value: string }).value))
    .toEqual(["A", "B", "C", "D", "E"]);
});

Assigning into out[i] rather than pushing is what makes that pass. It is a one-character difference in the implementation and a whole class of data-corruption bug in production.

Retrying only the rows that failed

The reason to classify errors is that the retry pass should be narrow. Re-running the whole batch to recover two rows pays for a hundred completions and, if any call has a side effect, performs ninety-eight of them twice. The retry takes the failed indices and nothing else.

1. Filter the results to !r.ok && r.retryable. A 400 for an oversized image will fail identically on retry; sending it again is pure cost.
2. Re-run runBatch over just those input rows, keeping a map from the position in the retry batch back to the original index.
3. Merge by original index. Assert in a test that a row which failed once and succeeded on retry ends up at its original position with ok: true, and that a row which failed twice is still reported rather than dropped in the merge.
4. Assert the second pass called the provider exactly as many times as there were retryable failures — two, not five. This is the assertion that keeps the retry narrow as the code is refactored.

If the per-row work writes anything — a row in a table, a file, a webhook — the retry needs an idempotency key derived from the input row rather than from the attempt, or the merge test passes while the database gets two copies. That is the same argument as checking a completion before you use it: the failure path is code, and code that only runs during an incident is code nobody has ever executed.