Tests

Name the behaviour, not the function

def test_scheduler():
    s = Scheduler(FakeModel(), FakeTokenizer(), max_batch=1)
    s.waiting.extend([req("a"), req("b")])
    s.admit_waiting()
    assert list(s.running) == ["a"]
    s.advance()
    assert s.running["a"].generated == [FakeModel.TOKEN]
    s.cancel("a")
    s.admit_waiting()
    assert list(s.running) == ["b"]

def test_request_waits_while_the_batch_is_full():
    scheduler = make_scheduler(max_batch=1)
    scheduler.waiting.extend([req("a"), req("b")])
    scheduler.admit_waiting()
    assert list(scheduler.running) == ["a"]


def test_each_step_adds_one_token_per_running_sequence():
    scheduler = make_scheduler()
    scheduler.waiting.extend([req("a"), req("b")])
    scheduler.admit_waiting()
    scheduler.advance()
    assert [len(s.generated) for s in scheduler.running.values()] == [1, 1]


def test_cancelling_a_request_frees_its_slot():
    scheduler = make_scheduler(max_batch=1)
    scheduler.waiting.extend([req("a"), req("b")])
    scheduler.admit_waiting()
    scheduler.cancel("a")
    scheduler.admit_waiting()
    assert list(scheduler.running) == ["b"]

Read only the names on the right and you know three promises the scheduler makes. When one fails, the name says which promise broke.

Guidelines

Arrange, act, assert

Each test sets up, does one thing and checks the result, in that order, with a blank line between the parts if it helps. No logic (loops, conditionals) in the test body. If you need a loop, use pytest.mark.parametrize.

Fakes over mocks

A small hand-written FakeModel with real behaviour (it returns a fixed token, or echoes the prompt) gives readable tests that survive refactors. Heavy mock.patch setups test the implementation's call sequence, and break whenever it changes.

Test the promises that matter most

For a batching generation server those are usually:

• Progress: every running sequence gains exactly one token per step.
• Order: waiting requests are admitted in arrival order.
• Limits: a sequence stops at end-of-sequence or at max_new_tokens, whichever comes first.
• Backpressure: a full queue rejects new requests cleanly instead of growing without bound.
• Batch invariance: a request's output doesn't depend on which other requests shared its batch (with a deterministic fake model).

Fast by default

The core logic should be testable on a laptop CPU in seconds, with no GPU, network or model weights. Mark the slow, GPU-only tests (@pytest.mark.gpu) and run them separately.

Helpers that build, not helpers that assert

A make_scheduler(max_batch=2) factory keeps tests short. A custom assert_scheduler_ok(s) hides what's being checked; inline the assertions.