test: use stream_infer instead of generate

Author: windreamer

tests/test_lmdeploy/test_grammar.py

@@ -1,12 +1,11 @@
-import asyncio
 import json
 import re
 
 import pytest
 from jsonschema import validate
 
 from lmdeploy import pipeline
-from lmdeploy.messages import GenerationConfig, PytorchEngineConfig, Response, TurbomindEngineConfig
+from lmdeploy.messages import GenerationConfig, PytorchEngineConfig, TurbomindEngineConfig
 
 MODEL_IDS = [
     'Qwen/Qwen3-0.6B',
@@ -98,28 +97,6 @@ def test_guided_matrix(model_id, backend_name, backend_factory, schema_type):
         pipe.close()
 
 
-async def collect(*aiters):
-    results = [[] for _ in range(len(aiters))]
-
-    async def drain(idx, aiter):
-        async for item in aiter:
-            results[idx].append(item)
-
-    await asyncio.gather(*(drain(idx, aiter) for idx, aiter in enumerate(aiters)))
-
-    responses = []
-    for r in results:
-        resp = Response(text='', input_token_len=0, generate_token_len=0)
-        responses.append(resp)
-        for out in r:
-            resp.text += out.response
-            resp.input_token_len = out.input_token_len
-            resp.generate_token_len = out.generate_token_len
-            resp.finish_reason = out.finish_reason
-
-    return responses
-
-
 @pytest.mark.parametrize('model_id', MODEL_IDS)
 @pytest.mark.parametrize('backend_name,backend_factory', BACKEND_FACTORIES)
 def test_mix_guided_matrix(model_id, backend_name, backend_factory):
@@ -134,17 +111,15 @@ def test_mix_guided_matrix(model_id, backend_name, backend_factory):
     schema = SCHEMA_MAP[schema_type]
     response_format[schema_type] = dict(name='test', schema=schema)
 
-    gen_config = GenerationConfig(response_format=response_format)
+    prompts = ['Make a self introduction please.'] * 4
+    config = GenerationConfig(response_format=response_format)
+
+    gen_config = [None if idx % 3 else config for idx in range(4)]
 
-    configs = [None if idx % 3 else gen_config for idx in range(4)]
-    tasks = [
-        pipe.generate(messages='Make a self introduction please.', session_id=session_id, gen_config=gen_config)
-        for session_id, gen_config in enumerate(configs)
-    ]
+    responses = pipe.stream_infer(prompts, gen_config=gen_config)
 
-    responses = asyncio.run(collect(*tasks))
-    for resp, config in zip(responses, configs):
-        if config is None:
+    for resp, c in zip(responses, gen_config):
+        if c is None:
             assert '}' not in resp.text
         else:
             validate(instance=json.loads(resp.text), schema=schema)