check env inside code string

gaogaotiantian · Jan 22, 2024 · 090a21a · 090a21a
1 parent f47abc4
commit 090a21a
Showing 1 changed file with 19 additions and 14 deletions.
diff --git a/tests/test_regression.py b/tests/test_regression.py
@@ -387,11 +387,15 @@ def test_escape_string(self):
 
 
 wait_for_child = """
+import os
 import time
 import multiprocessing
 
 def target():
-    time.sleep(%s)
+    if os.getenv("GITHUB_ACTIONS"):
+        time.sleep(9)
+    else:
+        time.sleep(3)
 
 if __name__ == '__main__':
     p = multiprocessing.Process(target=target)
@@ -400,7 +404,10 @@ def target():
     # This is a hack to make sure the main process won't join the child process,
     # so we can test the VizUI.wait_children_finish function
     multiprocessing.process._children = set()
-    time.sleep(%s)
+    if os.getenv("GITHUB_ACTIONS"):
+        time.sleep(3)
+    else:
+        time.sleep(1)
 """
 
 wait_for_terminated_child = """
@@ -410,7 +417,10 @@ def target():
 import multiprocessing
 
 def target():
-    time.sleep(%s)
+    if os.getenv("GITHUB_ACTIONS"):
+        time.sleep(9)
+    else:
+        time.sleep(3)
     os.kill(os.getpid(), signal.SIGTERM)
 
 if __name__ == '__main__':
@@ -420,25 +430,20 @@ def target():
     # This is a hack to make sure the main process won't join the child process,
     # so we can test the VizUI.wait_children_finish function
     multiprocessing.process._children = set()
-    time.sleep(%s)
+    if os.getenv("GITHUB_ACTIONS"):
+        time.sleep(3)
+    else:
+        time.sleep(1)
 """
 
 
 class TestWaitForChild(CmdlineTmpl):
     def test_child_process_exits_normally(self):
-        if os.getenv("GITHUB_ACTIONS"):
-            wait_for_child_script = wait_for_child % (9, 3)
-        else:
-            wait_for_child_script = wait_for_child % (3, 1)
         self.template(["viztracer", "-o", "result.json", "cmdline_test.py"],
                       expected_output_file="result.json", expected_stdout=r"Wait",
-                      script=wait_for_child_script)
+                      script=wait_for_child)
 
     def test_child_process_exits_abnormally(self):
-        if os.getenv("GITHUB_ACTIONS"):
-            wait_for_terminated_child_script = wait_for_terminated_child % (9, 3)
-        else:
-            wait_for_terminated_child_script = wait_for_terminated_child % (3, 1)
         self.template(["viztracer", "-o", "result.json", "cmdline_test.py"],
                       expected_output_file="result.json", expected_stdout=r"Wait",
-                      script=wait_for_terminated_child_script)
+                      script=wait_for_terminated_child)