diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py
index c26070ad4..f1cf4f9b7 100644
--- a/tests/diffusion/test_stage_diffusion_proc.py
+++ b/tests/diffusion/test_stage_diffusion_proc.py
@@ -24,19 +24,29 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
                 SimpleNamespace(
                     images=["img-1"],
                     _multimodal_output={},
+                    _custom_output={},
                     metrics={},
                     stage_durations={},
                     peak_memory_mb=0.0,
                     latents=None,
+                    trajectory_latents=None,
+                    trajectory_timesteps=None,
+                    trajectory_log_probs=None,
+                    trajectory_decoded=None,
                     final_output_type="image",
                 ),
                 SimpleNamespace(
                     images=["img-2"],
                     _multimodal_output={},
+                    _custom_output={},
                     metrics={},
                     stage_durations={},
                     peak_memory_mb=0.0,
                     latents=None,
+                    trajectory_latents=None,
+                    trajectory_timesteps=None,
+                    trajectory_log_probs=None,
+                    trajectory_decoded=None,
                     final_output_type="image",
                 ),
             ]