From 3823fc557dffdef05ac7eff6e5db65b204306f9c Mon Sep 17 00:00:00 2001
From: YWMditto <ditto@YWMdittodeMacBook-Pro.local>
Date: Wed, 20 Apr 2022 00:16:23 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AF=B9=20magic=5Fargv=5Fenv=5Fcontext=20?=
 =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E4=BA=86=20timeout=20=E5=8F=82=E6=95=B0?=
 =?UTF-8?q?=EF=BC=8C=E6=B5=8B=E8=AF=95=E5=87=BD=E6=95=B0=E8=B6=85=E8=BF=87?=
 =?UTF-8?q?=E4=B8=80=E5=AE=9A=E6=97=B6=E9=97=B4=E5=90=8E=E8=87=AA=E5=8A=A8?=
 =?UTF-8?q?kill=E6=8E=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/helpers/utils.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/tests/helpers/utils.py b/tests/helpers/utils.py
index 9a4af07c..c0b51a8b 100644
--- a/tests/helpers/utils.py
+++ b/tests/helpers/utils.py
@@ -1,11 +1,12 @@
 import os
 import sys
 import __main__
-from functools import wraps
+from functools import wraps, partial
 from inspect import ismethod
 from copy import deepcopy
 from io import StringIO
 import time
+import signal
 
 import numpy as np
 
@@ -29,7 +30,15 @@ def recover_logger(fn):
     return wrapper
 
 
-def magic_argv_env_context(fn):
+def magic_argv_env_context(fn=None, timeout=600):
+    """
+    用来在测试时包裹每一个单独的测试函数，使得 ddp 测试正确；
+    :param timeout: 表示一个测试如果经过多久还没有通过的话就主动将其 kill 掉，默认为 10 分钟，单位为秒；
+    :return:
+    """
+    # 说明是通过 @magic_argv_env_context(timeout=600) 调用；
+    if fn is None:
+        return partial(magic_argv_env_context, timeout=timeout)
 
     @wraps(fn)
     def wrapper(*args, **kwargs):
@@ -55,11 +64,17 @@ def magic_argv_env_context(fn):
         else:
             sys.argv = [sys.argv[0], f"{os.path.abspath(sys.modules[fn.__module__].__file__)}::{get_class_that_defined_method(fn).__name__}::{fn.__name__}{subtest}"] + used_args
 
+        def _handle_timeout(signum, frame):
+            raise TimeoutError(f"\nYour test fn: {fn.__name__} has timed out.\n")
+
+        signal.signal(signal.SIGALRM, _handle_timeout)
+        signal.alarm(timeout)
         res = fn(*args, **kwargs)
+        signal.alarm(0)
         sys.argv = deepcopy(command)
         os.environ = env
-
         return res
+
     return wrapper