diff --git a/python/sglang/multimodal_gen/test/unit/test_glm_image_ar.py b/python/sglang/multimodal_gen/test/unit/test_glm_image_ar.py index 5c1f89545..3661df8f1 100644 --- a/python/sglang/multimodal_gen/test/unit/test_glm_image_ar.py +++ b/python/sglang/multimodal_gen/test/unit/test_glm_image_ar.py @@ -47,6 +47,17 @@ class _FakeResponse: return data +class _FakeBatchResponse: + def __init__(self, outputs): + self._outputs = outputs + + def raise_for_status(self): + return None + + def json(self): + return self._outputs + + class TestGlmImageARSrtBackend(unittest.TestCase): def _server_args(self): return SimpleNamespace( @@ -135,16 +146,18 @@ class TestGlmImageARSrtBackend(unittest.TestCase): ) def test_srt_ar_forward_aggregates_usage(self, mock_post, _mock_device): set_global_server_args(self._server_args()) - mock_post.side_effect = [ - _FakeResponse( - list(range(1025)), - meta_info={"prompt_tokens": 13, "completion_tokens": 25}, - ), - _FakeResponse( - list(range(1025)), - meta_info={"prompt_tokens": 13, "completion_tokens": 25}, - ), - ] + mock_post.return_value = _FakeBatchResponse( + [ + { + "output_ids": list(range(1025)), + "meta_info": {"prompt_tokens": 13, "completion_tokens": 25}, + }, + { + "output_ids": list(range(1025)), + "meta_info": {"prompt_tokens": 13, "completion_tokens": 25}, + }, + ] + ) stage = GlmImageAR(processor=_FakeProcessor(), vision_language_encoder=None) batch = SimpleNamespace( prompt="A simple product sketch", @@ -153,6 +166,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase): image_path=None, num_outputs_per_prompt=2, seed=None, + extra={}, ) batch = stage.forward(batch, self._server_args()) diff --git a/python/sglang/multimodal_gen/test/unit/test_glm_image_multi_output.py b/python/sglang/multimodal_gen/test/unit/test_glm_image_multi_output.py index ee797e84b..ec87062bc 100644 --- a/python/sglang/multimodal_gen/test/unit/test_glm_image_multi_output.py +++ b/python/sglang/multimodal_gen/test/unit/test_glm_image_multi_output.py @@ -30,7 +30,7 @@ class _RecordingGlmImageAR(GlmImageAR): def generate_prior_tokens(self, **kwargs): self.initial_seeds.append(torch.initial_seed()) output_idx = len(self.initial_seeds) - return torch.full((1, 4), output_idx, dtype=torch.long), None + return torch.full((1, 4), output_idx, dtype=torch.long), None, None class _DummySchedulerConfig(dict): @@ -105,6 +105,7 @@ def test_ar_stage_generates_one_prior_per_requested_output(): image_path=None, num_outputs_per_prompt=2, seed=11, + extra={}, ) with patch.object(