[CI] Fix GLM-Image usage unit tests (#33731)

This commit is contained in:
Mohammad Miadh Angkad
2026-08-05 16:20:08 -07:00
committed by GitHub
parent 070fde72bf
commit 5424d2039c
2 changed files with 26 additions and 11 deletions
@@ -47,6 +47,17 @@ class _FakeResponse:
return data 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): class TestGlmImageARSrtBackend(unittest.TestCase):
def _server_args(self): def _server_args(self):
return SimpleNamespace( return SimpleNamespace(
@@ -135,16 +146,18 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
) )
def test_srt_ar_forward_aggregates_usage(self, mock_post, _mock_device): def test_srt_ar_forward_aggregates_usage(self, mock_post, _mock_device):
set_global_server_args(self._server_args()) set_global_server_args(self._server_args())
mock_post.side_effect = [ mock_post.return_value = _FakeBatchResponse(
_FakeResponse( [
list(range(1025)), {
meta_info={"prompt_tokens": 13, "completion_tokens": 25}, "output_ids": list(range(1025)),
), "meta_info": {"prompt_tokens": 13, "completion_tokens": 25},
_FakeResponse( },
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) stage = GlmImageAR(processor=_FakeProcessor(), vision_language_encoder=None)
batch = SimpleNamespace( batch = SimpleNamespace(
prompt="A simple product sketch", prompt="A simple product sketch",
@@ -153,6 +166,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
image_path=None, image_path=None,
num_outputs_per_prompt=2, num_outputs_per_prompt=2,
seed=None, seed=None,
extra={},
) )
batch = stage.forward(batch, self._server_args()) batch = stage.forward(batch, self._server_args())
@@ -30,7 +30,7 @@ class _RecordingGlmImageAR(GlmImageAR):
def generate_prior_tokens(self, **kwargs): def generate_prior_tokens(self, **kwargs):
self.initial_seeds.append(torch.initial_seed()) self.initial_seeds.append(torch.initial_seed())
output_idx = len(self.initial_seeds) 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): class _DummySchedulerConfig(dict):
@@ -105,6 +105,7 @@ def test_ar_stage_generates_one_prior_per_requested_output():
image_path=None, image_path=None,
num_outputs_per_prompt=2, num_outputs_per_prompt=2,
seed=11, seed=11,
extra={},
) )
with patch.object( with patch.object(