From aefa2ee45cd0f7145436e265596bb22853e18241 Mon Sep 17 00:00:00 2001
From: Jinheng <ahengljh@gmail.com>
Date: Wed, 8 Apr 2026 17:45:10 +0800
Subject: [PATCH] [Fix] Align diffusion proc test mock with current output
 fields (#2584)

Signed-off-by: Jinheng Li <ahengljh@gmail.com>
---
 tests/diffusion/test_stage_diffusion_proc.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

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",
                 ),
             ]