2
2
import asyncio
3
3
import os
4
4
import sys
5
+ import shutil
5
6
import tempfile
6
7
from pathlib import Path
7
8
from typing import AsyncGenerator , TypeAlias
@@ -32,7 +33,7 @@ def docker_tests_enabled() -> bool:
32
33
return False
33
34
34
35
35
- @pytest_asyncio .fixture (scope = "function " ) # type: ignore
36
+ @pytest_asyncio .fixture (scope = "module " ) # type: ignore
36
37
async def executor_and_temp_dir (
37
38
request : pytest .FixtureRequest ,
38
39
) -> AsyncGenerator [tuple [DockerCommandLineCodeExecutor , str ], None ]:
@@ -47,9 +48,20 @@ async def executor_and_temp_dir(
47
48
ExecutorFixture : TypeAlias = tuple [DockerCommandLineCodeExecutor , str ]
48
49
49
50
51
+ @pytest_asyncio .fixture (scope = "function" ) # type: ignore
52
+ async def cleanup_temp_dir (executor_and_temp_dir : ExecutorFixture ) -> AsyncGenerator [None , None ]:
53
+ _executor , temp_dir = executor_and_temp_dir
54
+ for file in Path (temp_dir ).iterdir ():
55
+ if file .is_file ():
56
+ file .unlink ()
57
+ elif file .is_dir ():
58
+ shutil .rmtree (file )
59
+ yield None
60
+
61
+
50
62
@pytest .mark .asyncio
51
63
@pytest .mark .parametrize ("executor_and_temp_dir" , ["docker" ], indirect = True )
52
- async def test_execute_code (executor_and_temp_dir : ExecutorFixture ) -> None :
64
+ async def test_execute_code (executor_and_temp_dir : ExecutorFixture , cleanup_temp_dir : None ) -> None :
53
65
executor , _temp_dir = executor_and_temp_dir
54
66
cancellation_token = CancellationToken ()
55
67
@@ -97,7 +109,9 @@ async def test_execute_code(executor_and_temp_dir: ExecutorFixture) -> None:
97
109
98
110
@pytest .mark .asyncio
99
111
@pytest .mark .parametrize ("executor_and_temp_dir" , ["docker" ], indirect = True )
100
- async def test_commandline_code_executor_timeout (executor_and_temp_dir : ExecutorFixture ) -> None :
112
+ async def test_commandline_code_executor_timeout (
113
+ executor_and_temp_dir : ExecutorFixture , cleanup_temp_dir : None
114
+ ) -> None :
101
115
_executor , temp_dir = executor_and_temp_dir
102
116
cancellation_token = CancellationToken ()
103
117
code_blocks = [CodeBlock (code = "import time; time.sleep(10); print('hello world!')" , language = "python" )]
@@ -110,7 +124,9 @@ async def test_commandline_code_executor_timeout(executor_and_temp_dir: Executor
110
124
111
125
@pytest .mark .asyncio
112
126
@pytest .mark .parametrize ("executor_and_temp_dir" , ["docker" ], indirect = True )
113
- async def test_commandline_code_executor_cancellation (executor_and_temp_dir : ExecutorFixture ) -> None :
127
+ async def test_commandline_code_executor_cancellation (
128
+ executor_and_temp_dir : ExecutorFixture , cleanup_temp_dir : None
129
+ ) -> None :
114
130
_executor , temp_dir = executor_and_temp_dir
115
131
cancellation_token = CancellationToken ()
116
132
# Write code that sleep for 10 seconds and then write "hello world!"
@@ -137,7 +153,7 @@ async def test_commandline_code_executor_cancellation(executor_and_temp_dir: Exe
137
153
138
154
@pytest .mark .asyncio
139
155
@pytest .mark .parametrize ("executor_and_temp_dir" , ["docker" ], indirect = True )
140
- async def test_invalid_relative_path (executor_and_temp_dir : ExecutorFixture ) -> None :
156
+ async def test_invalid_relative_path (executor_and_temp_dir : ExecutorFixture , cleanup_temp_dir : None ) -> None :
141
157
executor , _temp_dir = executor_and_temp_dir
142
158
cancellation_token = CancellationToken ()
143
159
code = """# filename: /tmp/test.py
@@ -152,7 +168,7 @@ async def test_invalid_relative_path(executor_and_temp_dir: ExecutorFixture) ->
152
168
153
169
@pytest .mark .asyncio
154
170
@pytest .mark .parametrize ("executor_and_temp_dir" , ["docker" ], indirect = True )
155
- async def test_valid_relative_path (executor_and_temp_dir : ExecutorFixture ) -> None :
171
+ async def test_valid_relative_path (executor_and_temp_dir : ExecutorFixture , cleanup_temp_dir : None ) -> None :
156
172
executor , temp_dir_str = executor_and_temp_dir
157
173
158
174
cancellation_token = CancellationToken ()
@@ -243,18 +259,13 @@ async def test_docker_commandline_code_executor_extra_args() -> None:
243
259
async def test_docker_commandline_code_executor_serialization () -> None :
244
260
with tempfile .TemporaryDirectory () as temp_dir :
245
261
executor = DockerCommandLineCodeExecutor (work_dir = temp_dir )
246
- await executor .start ()
247
262
248
263
executor_config = executor .dump_component ()
249
264
loaded_executor = DockerCommandLineCodeExecutor .load_component (executor_config )
250
- await loaded_executor .start ()
251
265
252
266
assert executor .bind_dir == loaded_executor .bind_dir
253
267
assert executor .timeout == loaded_executor .timeout
254
268
255
- await executor .stop ()
256
- await loaded_executor .stop ()
257
-
258
269
259
270
def test_invalid_timeout () -> None :
260
271
with pytest .raises (ValueError , match = "Timeout must be greater than or equal to 1." ):
@@ -269,23 +280,23 @@ async def test_directory_not_initialized() -> None:
269
280
270
281
271
282
@pytest .mark .asyncio
272
- async def test_error_wrong_path () -> None :
283
+ @pytest .mark .parametrize ("executor_and_temp_dir" , ["docker" ], indirect = True )
284
+ async def test_error_wrong_path (executor_and_temp_dir : ExecutorFixture , cleanup_temp_dir : None ) -> None :
273
285
if not docker_tests_enabled ():
274
286
pytest .skip ("Docker tests are disabled" )
275
287
276
- with tempfile .TemporaryDirectory () as temp_dir :
277
- async with DockerCommandLineCodeExecutor (work_dir = temp_dir ) as executor :
278
- cancellation_token = CancellationToken ()
279
- code_blocks = [
280
- CodeBlock (
281
- code = """with open("/nonexistent_dir/test.txt", "w") as f:
282
- f.write("hello world!")""" ,
283
- language = "python" ,
284
- )
285
- ]
286
- result = await executor .execute_code_blocks (code_blocks , cancellation_token )
287
- assert result .exit_code != 0
288
- assert "No such file or directory" in result .output
288
+ executor , _ = executor_and_temp_dir
289
+ cancellation_token = CancellationToken ()
290
+ code_blocks = [
291
+ CodeBlock (
292
+ code = """with open("/nonexistent_dir/test.txt", "w") as f:
293
+ f.write("hello world!")""" ,
294
+ language = "python" ,
295
+ )
296
+ ]
297
+ result = await executor .execute_code_blocks (code_blocks , cancellation_token )
298
+ assert result .exit_code != 0
299
+ assert "No such file or directory" in result .output
289
300
290
301
291
302
@pytest .mark .asyncio
@@ -364,7 +375,10 @@ async def test_delete_tmp_files() -> None:
364
375
365
376
366
377
@pytest .mark .asyncio
367
- async def test_docker_commandline_code_executor_with_multiple_tasks () -> None :
378
+ @pytest .mark .parametrize ("executor_and_temp_dir" , ["docker" ], indirect = True )
379
+ async def test_docker_commandline_code_executor_with_multiple_tasks (
380
+ executor_and_temp_dir : ExecutorFixture , cleanup_temp_dir : None
381
+ ) -> None :
368
382
if not docker_tests_enabled ():
369
383
pytest .skip ("Docker tests are disabled" )
370
384
@@ -382,6 +396,5 @@ async def run_cancellation_scenario(executor: DockerCommandLineCodeExecutor) ->
382
396
def run_scenario_in_new_loop (executor_instance : DockerCommandLineCodeExecutor ) -> None :
383
397
asyncio .run (run_cancellation_scenario (executor_instance ))
384
398
385
- with tempfile .TemporaryDirectory () as temp_dir :
386
- async with DockerCommandLineCodeExecutor (work_dir = temp_dir ) as executor :
387
- await asyncio .get_running_loop ().run_in_executor (None , run_scenario_in_new_loop , executor )
399
+ executor , _ = executor_and_temp_dir
400
+ await asyncio .get_running_loop ().run_in_executor (None , run_scenario_in_new_loop , executor )
0 commit comments