|
6 | 6 |
|
7 | 7 | from __future__ import annotations |
8 | 8 | from typing_extensions import TYPE_CHECKING |
9 | | -from concurrent.futures import ThreadPoolExecutor |
| 9 | +from concurrent.futures import Executor, ThreadPoolExecutor |
10 | 10 | from os import environ |
11 | 11 | from pathlib import Path |
12 | 12 |
|
|
24 | 24 | from .deepequals import assert_deep_equal, assert_not_deep_equal |
25 | 25 |
|
26 | 26 | # 并行 workers 数 |
27 | | -MAX_WORKERS = 2 |
| 27 | +MAX_WORKERS = int(environ.get("MAX_WORKERS", "2")) |
28 | 28 |
|
29 | 29 | # 跳过大型文件 |
30 | 30 | SKIP_HUGE = environ.get("SKIP_HUGE", "0") != "0" |
@@ -127,22 +127,14 @@ def test_mira_file(mira_file: Path) -> None: |
127 | 127 | if is_huge and SKIP_HUGE: |
128 | 128 | pytest.skip("Skipping huge test file") |
129 | 129 |
|
130 | | - pool = ThreadPoolExecutor(max_workers=MAX_WORKERS) if MAX_WORKERS > 0 else None |
131 | | - |
132 | | - try: |
133 | | - timeout_fns = None |
134 | | - config_checkpoint(120 if is_huge else 10) |
135 | | - if pool is not None: |
136 | | - futures = [ |
137 | | - pool.submit(_run_mira_file, mira_file) for _ in range(MAX_WORKERS) |
138 | | - ] |
139 | | - timeout_fns = _run_mira_file(mira_file) |
140 | | - for future in futures: |
141 | | - future.result() |
142 | | - else: |
143 | | - timeout_fns = _run_mira_file(mira_file) |
| 130 | + with ( |
| 131 | + ThreadPoolExecutor(max_workers=MAX_WORKERS) if MAX_WORKERS > 0 else Executor() |
| 132 | + ) as pool: |
| 133 | + config_checkpoint(300 if is_huge else 10) |
| 134 | + |
| 135 | + for _ in range(MAX_WORKERS): |
| 136 | + pool.submit(_run_mira_file, mira_file) |
| 137 | + |
| 138 | + timeout_fns = _run_mira_file(mira_file) |
144 | 139 | config_checkpoint() # 重置检查点配置 |
145 | 140 | _run_timeout_fns(timeout_fns) |
146 | | - finally: |
147 | | - if pool is not None: |
148 | | - pool.shutdown(wait=False) |
0 commit comments