diff --git a/tests/policies/pi0_fast/test_pi0_fast_parity.py b/tests/policies/pi0_fast/test_pi0_fast_parity.py
index f9e1a3b12..b955742c8 100644
--- a/tests/policies/pi0_fast/test_pi0_fast_parity.py
+++ b/tests/policies/pi0_fast/test_pi0_fast_parity.py
@@ -62,10 +62,11 @@ def test_pi0_fast_action_tokens_have_no_second_bos():
     step.action_tokenizer = lambda _: torch.tensor([4, 9])
     step._paligemma_tokenizer = _FakePaliGemmaTokenizer()
 
-    tokens, mask = step._tokenize_action(torch.zeros(1, 2, 1))
+    tokens, mask, code_mask = step._tokenize_action(torch.zeros(1, 2, 1))
 
     mapped = [1000 - 1 - 128 - token for token in (4, 9)]
     assert tokens[0, :6].tolist() == [10, 11, *mapped, 12, 1]
+    assert code_mask[0].tolist() == [False, False, True, True, False, False, False, False]
     assert _FakePaliGemmaTokenizer.bos_token_id not in tokens[0, mask[0]].tolist()
 
 
@@ -84,6 +85,7 @@ def test_action_tokenizer_serializes_bos_layout():
         "max_action_tokens": 8,
         "fast_skip_tokens": 128,
         "paligemma_tokenizer_name": "paligemma",
+        "allow_truncation": True,
         "prepend_bos": False,
         "action_tokenizer_name": "fast",
     }
diff --git a/tests/processor/test_tokenizer_processor.py b/tests/processor/test_tokenizer_processor.py
index aab0c14ec..cb9a09ae4 100644
--- a/tests/processor/test_tokenizer_processor.py
+++ b/tests/processor/test_tokenizer_processor.py
@@ -104,6 +104,7 @@ def test_action_tokenizer_config_preserves_token_mapping():
         "fast_skip_tokens": 64,
         "paligemma_tokenizer_name": "custom/paligemma",
         "allow_truncation": False,
+        "prepend_bos": True,
         "action_tokenizer_name": "custom/fast",
     }