@@ -159,28 +159,45 @@ async def test_repeated_submit_same_iens(driver: Driver, tmp_path, monkeypatch):
159159@pytest .mark .flaky (reruns = 5 )
160160async def test_kill_actually_kills (driver : Driver , tmp_path , pytestconfig , monkeypatch ):
161161 monkeypatch .chdir (tmp_path )
162- finished = False
162+ realizations = {0 , 1 }
163+ iens_status = {
164+ realization : {"finished" : False , "started" : False }
165+ for realization in realizations
166+ }
167+
163168 driver ._poll_period = 0.01
164169
165- async def kill_job_once_started (iens ):
166- nonlocal driver
167- await driver .kill (iens )
170+ async def kill_jobs_once_all_started (realizations ):
171+ nonlocal driver , iens_status
172+ for realization in realizations :
173+ iens_status [realization ]["started" ] = True
174+ if all (realization ["started" ] for realization in iens_status .values ()):
175+ await driver .kill (iens_status .keys ())
168176
169177 async def mark_as_finished (iens , code ):
170- nonlocal finished
171- finished = True
172-
173- await driver .submit (
174- 0 ,
175- "sh" ,
176- "-c" ,
177- f"sleep 10; touch { tmp_path } /survived" ,
178- name = "kill_me" ,
178+ nonlocal iens_status
179+ iens_status [iens ]["finished" ] = True
180+
181+ for realization_id in realizations :
182+ await driver .submit (
183+ realization_id ,
184+ "sh" ,
185+ "-c" ,
186+ f"sleep 10; touch { tmp_path } /{ realization_id } " ,
187+ name = f"kill_me_{ realization_id } " ,
188+ )
189+
190+ await poll (
191+ driver ,
192+ realizations ,
193+ started = kill_jobs_once_all_started ,
194+ finished = mark_as_finished ,
179195 )
180- await poll (driver , {0 }, started = kill_job_once_started , finished = mark_as_finished )
181- assert finished
182-
183- assert not Path ("survived" ).exists (), "Job should have been killed"
196+ for realization_id in realizations :
197+ assert iens_status [realization_id ]["finished" ]
198+ assert not Path (f"survived_{ realization_id } " ).exists (), (
199+ "Job should have been killed"
200+ )
184201
185202
186203@pytest .mark .integration_test
0 commit comments