44import asyncio
55import base64
66import json
7+ import warnings
78
89import librosa
910import numpy as np
@@ -85,7 +86,41 @@ async def test_multi_chunk_streaming(
8586
8687 await send_event (ws , {"type" : "session.update" , "model" : model_name })
8788
88- # Send commit to start transcription
89+ # Wait for the server to acknowledge the session update.
90+ try :
91+ while True :
92+ event = await receive_event (ws , timeout = 5.0 )
93+ if event ["type" ] == "session.updated" :
94+ break
95+ except TimeoutError :
96+ warnings .warn (
97+ f"session.updated not received within { 5.0 } s after "
98+ "session.update. The server may not implement this event." ,
99+ stacklevel = 2 ,
100+ )
101+
102+ # (ROCm) Warm-up: send a non-final commit (required to start
103+ # transcription) with a small audio chunk to trigger aiter
104+ # compilation on first use.
105+ await send_event (ws , {"type" : "input_audio_buffer.commit" })
106+ await send_event (
107+ ws ,
108+ {
109+ "type" : "input_audio_buffer.append" ,
110+ "audio" : mary_had_lamb_audio_chunks [0 ],
111+ },
112+ )
113+ await send_event (ws , {"type" : "input_audio_buffer.commit" , "final" : True })
114+
115+ # (ROCm) Drain all warm-up responses with generous timeout for
116+ # JIT compilation
117+ warmup_done = False
118+ while not warmup_done :
119+ event = await receive_event (ws , timeout = 360.0 )
120+ if event ["type" ] in ("transcription.done" , "error" ):
121+ warmup_done = True
122+
123+ # Now send the real test audio
89124 await send_event (ws , {"type" : "input_audio_buffer.commit" })
90125
91126 # Send multiple audio chunks
@@ -153,6 +188,18 @@ async def test_empty_commit_does_not_crash_engine(
153188
154189 await send_event (ws , {"type" : "session.update" , "model" : model_name })
155190
191+ try :
192+ while True :
193+ event = await receive_event (ws , timeout = 5.0 )
194+ if event ["type" ] == "session.updated" :
195+ break
196+ except TimeoutError :
197+ warnings .warn (
198+ f"session.updated not received within { 5.0 } s after "
199+ "session.update. The server may not implement this event." ,
200+ stacklevel = 2 ,
201+ )
202+
156203 # Start generation without sending any audio
157204 await send_event (ws , {"type" : "input_audio_buffer.commit" })
158205
@@ -161,7 +208,8 @@ async def test_empty_commit_does_not_crash_engine(
161208
162209 # We should get *some* response (error or empty transcription),
163210 # but the engine must NOT crash.
164- event = await receive_event (ws , timeout = 30.0 )
211+ # (ROCm) Use generous timeout for first request (aiter JIT compilation)
212+ event = await receive_event (ws , timeout = 360.0 )
165213 assert event ["type" ] in (
166214 "error" ,
167215 "transcription.done" ,
@@ -176,6 +224,19 @@ async def test_empty_commit_does_not_crash_engine(
176224
177225 await send_event (ws , {"type" : "session.update" , "model" : model_name })
178226
227+ try :
228+ while True :
229+ event = await receive_event (ws , timeout = 5.0 )
230+ if event ["type" ] == "session.updated" :
231+ break
232+ except TimeoutError :
233+ warnings .warn (
234+ f"session.updated not received within { 5.0 } s after "
235+ "session.update. The server may not implement this event." ,
236+ stacklevel = 2 ,
237+ )
238+
239+ # Start transcription
179240 await send_event (ws , {"type" : "input_audio_buffer.commit" })
180241
181242 for chunk in mary_had_lamb_audio_chunks :
0 commit comments