diff --git a/tests/test_multiprocess.py b/tests/test_multiprocess.py index dfb8ce14..543344f8 100644 --- a/tests/test_multiprocess.py +++ b/tests/test_multiprocess.py @@ -45,6 +45,12 @@ def fib(n): print(subprocess.call(["python", "-m", "timeit", "-n", "100", "'1+1'"])) """ +file_subprocess_code_string = """ +import subprocess +p = subprocess.Popen(['python', '-c', 'import time;time.sleep(0.5)']) +p.wait() +""" + file_fork = """ import os import time @@ -227,6 +233,13 @@ def setUp(self): def tearDown(self): os.remove("child.py") + def assertSubprocessName(self, name, data): + for entry in data["traceEvents"]: + if entry["name"] == "process_name" and entry["args"]["name"] == name: + break + else: + self.fail("no matching subprocess name") + def test_basic(self): def check_func(data): pids = set() @@ -242,13 +255,7 @@ def test_child_process(self): expected_output_file=None) self.assertEqual(len(os.listdir(tmpdir)), 1) with open(os.path.join(tmpdir, os.listdir(tmpdir)[0])) as f: - trace_events = json.load(f)["traceEvents"] - for entry in trace_events: - if entry["name"] == "process_name": - self.assertNotEqual(entry["args"]["name"], "MainProcess") - break - else: - self.fail("no process_name event found") + self.assertSubprocessName("child.py", json.load(f)) def test_module(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -256,28 +263,23 @@ def test_module(self): self.template(["viztracer", "-o", output_file, "cmdline_test.py"], expected_output_file=output_file, expected_stdout=".*100 loops.*", - script=file_subprocess_module, cleanup=False) - with open(output_file) as f: - trace_events = json.load(f)["traceEvents"] - for entry in trace_events: - if entry["name"] == "process_name" and entry["args"]["name"] == "timeit": - break - else: - self.fail("no valid child process_name event found") - - def test_code(self): + script=file_subprocess_module, + check_func=lambda data: self.assertSubprocessName("timeit", data)) + + def test_code_string(self): with tempfile.TemporaryDirectory() as tmpdir: + output_path = os.path.join(tmpdir, "result.json") + self.template(["viztracer", "-o", output_path, "cmdline_test.py"], + expected_output_file=output_path, + script=file_subprocess_code_string, + check_func=lambda data: self.assertSubprocessName("python -c", data)) + + # this is for coverage self.template(['viztracer', '-o', os.path.join(tmpdir, "result.json"), '--subprocess_child', '-c', 'import time;time.sleep(0.5)'], expected_output_file=None) self.assertEqual(len(os.listdir(tmpdir)), 1) with open(os.path.join(tmpdir, os.listdir(tmpdir)[0])) as f: - trace_events = json.load(f)["traceEvents"] - for entry in trace_events: - if entry["name"] == "process_name": - self.assertEqual(entry["args"]["name"], "python -c") - break - else: - self.fail("no valid child process_name event found") + self.assertSubprocessName("python -c", json.load(f)) def test_nested(self): def check_func(data):