Why

Two small helpers in the package have public-API contracts that recently had docs/fix PRs land against them:

• basic_collation_fn in src/alpamayo_r1/processor/qwen_processor.py (the mutable-default-arg fix proposed in Replace mutable default list with None sentinel in basic_collation_fn #84).
• get_assistant_mask in src/alpamayo_r1/utils/get_label_mask.py (the return-type/docstring fix proposed in Fix get_label_mask docstrings and get_assistant_mask return type #85).

Without tests, a future edit can quietly regress either contract. This PR follows the same pattern as the to_device tests in #80 — add focused, dependency-light pytest cases that lock in the documented behavior.

What

src/alpamayo_r1/processor/test_qwen_processor.py — 5 tests covering basic_collation_fn:

• Default arg behaves as "no extra unstackable keys".
• unstackable_keys=None is accepted and equivalent to omitting it.
• Keys named in unstackable_keys come back as a plain list.
• Repeated calls don't accumulate state via the default arg.
• Non-tensor values are always returned as lists (never stacked).

src/alpamayo_r1/utils/test_get_label_mask.py — 4 tests covering get_assistant_mask:

• torch.Tensor input → torch.Tensor (bool) output.
• list[int] input → list[bool] output (the dual return shape the annotation now reflects).
• The mask covers exactly the assistant content + trailing EOS, with the user turn untouched (matches the body's start+3 / end+1 logic).
• Tensor and list input paths produce identical booleans.

Tests use a tiny stub tokenizer that only implements convert_tokens_to_ids, so they run on CI without HF auth or model downloads.

Run

pytest src/alpamayo_r1/processor/test_qwen_processor.py src/alpamayo_r1/utils/test_get_label_mask.py -v

Notes

• These tests pass on the current main for the get_assistant_mask file (the fix in Fix get_label_mask docstrings and get_assistant_mask return type #85 was docstring/annotation only; the runtime behavior these tests pin down already exists).
• The collation tests pass on current main too — they will also pass after Replace mutable default list with None sentinel in basic_collation_fn #84 lands, which is the point of a regression test.