From 747ebadd1b839dff7557125c13f5d49f820050b1 Mon Sep 17 00:00:00 2001 From: Tian Gao Date: Tue, 30 Jan 2024 12:52:51 -0800 Subject: [PATCH] Fix remote stuff --- src/viztracer/main.py | 33 +++++++++++++++++++++++++-------- tests/test_remote.py | 24 +++++++++++++++--------- 2 files changed, 40 insertions(+), 17 deletions(-) diff --git a/src/viztracer/main.py b/src/viztracer/main.py index 770b94e9..1b54222c 100644 --- a/src/viztracer/main.py +++ b/src/viztracer/main.py @@ -9,6 +9,7 @@ import json import multiprocessing.util # type: ignore import os +import platform import shutil import signal import sys @@ -479,15 +480,26 @@ def show_version(self) -> VizProcedureResult: print(__version__) return True, None + def _check_attach_availability(self) -> Tuple[bool, Optional[str]]: + if sys.platform == "win32": + return False, "VizTracer does not support this feature on Windows" + + if sys.platform == "darwin" and platform.processor() == "arm": + return False, "VizTracer does not support this feature on Apple Silicon" + + if sys.platform == "darwin" and sys.version_info >= (3, 11): + color_print("WARNING", "Warning: attach may not work on 3.11+ on Mac due to hardened runtime") + + return True, None + def attach(self) -> VizProcedureResult: pid = self.options.attach interval = self.options.t - if sys.platform == "win32": - return False, "VizTracer does not support this feature on Windows" + success, err_msg = self._check_attach_availability() - if sys.platform == "darwin" and sys.version_info >= (3, 11): - print("Warning: attach may not work on 3.11+ on Mac due to hardened runtime") + if not success: + return False, err_msg if not pid_exists(pid): return False, f"pid {pid} does not exist!" @@ -522,8 +534,10 @@ def attach(self) -> VizProcedureResult: def uninstall(self) -> VizProcedureResult: pid = self.options.uninstall - if sys.platform == "win32": - return False, "VizTracer does not support this feature on Windows" + success, err_msg = self._check_attach_availability() + + if not success: + return False, err_msg if not pid_exists(pid): return False, f"pid {pid} does not exist!" @@ -537,8 +551,11 @@ def uninstall(self) -> VizProcedureResult: return True, None def attach_installed(self) -> VizProcedureResult: - if sys.platform == "win32": - return False, "VizTracer does not support this feature on Windows" + success, err_msg = self._check_attach_availability() + + if not success: + return False, err_msg + pid = self.options.attach_installed interval = self.options.t try: diff --git a/tests/test_remote.py b/tests/test_remote.py index 74865345..ff271d8e 100644 --- a/tests/test_remote.py +++ b/tests/test_remote.py @@ -5,6 +5,7 @@ import base64 import json import os +import platform import re import signal import subprocess @@ -21,9 +22,14 @@ from .util import cmd_with_coverage -@unittest.skipIf(sys.platform == "darwin" and sys.version_info >= (3, 11), "Does not support 3.11+ on Mac") +attach_unavailable = (sys.platform == "win32" + or (sys.platform == "darwin" + and (sys.version_info > (3, 11) + or "arm" in platform.processor()))) + + +@unittest.skipIf(attach_unavailable, "Does not support attach on this platform") class TestRemote(CmdlineTmpl): - @unittest.skipIf(sys.platform == "win32", "Does not support on Windows") def test_install(self): tracer = VizTracer(output_file="remote.json", verbose=0) tracer.install() @@ -33,7 +39,6 @@ def test_install(self): self.assertFileExists("remote.json") os.remove("remote.json") - @unittest.skipIf(sys.platform == "win32", "Does not support on Windows") def test_attach_installed(self): file_to_attach = textwrap.dedent(""" from viztracer import VizTracer @@ -52,7 +57,6 @@ def test_attach_installed(self): self.attach_check(file_to_attach, attach_cmd, output_file) self.attach_check(file_to_attach, attach_installed_cmd, output_file, use_installed=True) - @unittest.skipIf(sys.platform == "win32", "Does not support on Windows") def test_attach(self): file_to_attach = textwrap.dedent(""" import time @@ -138,7 +142,6 @@ def attach_check(self, file_to_attach, attach_cmd, output_file, file_should_exis p_attach_invalid.wait() self.assertTrue(p_attach_invalid.returncode != 0) - @unittest.skipIf(sys.platform == "win32", "Does not support on Windows") def test_uninstall(self): file_to_attach = textwrap.dedent(""" import time @@ -208,7 +211,11 @@ def test_uninstall(self): p_attach_uninstall.wait() self.assertTrue(p_attach_uninstall.returncode != 0) - @unittest.skipIf(sys.platform != "win32", "Only test Windows") + +class TestRemoteFail(CmdlineTmpl): + @unittest.skipUnless(sys.platform == "win32" + or (sys.platform == "darwin" and "arm" in platform.processor()), + "Only test unavailable platform") def test_windows(self): tracer = VizTracer(output_file="remote.json") with self.assertRaises(SystemExit): @@ -219,9 +226,8 @@ def test_windows(self): self.template(["viztracer", "--uninstall", "1234"], success=False) -@unittest.skipIf(sys.platform == "darwin" and sys.version_info >= (3, 11), "Does not support 3.11+ on Mac") +@unittest.skipIf(attach_unavailable, "Does not support this platform") class TestAttachSanity(CmdlineTmpl): - @unittest.skipIf(sys.platform == "win32", "Can't run attach on Windows") def test_basic(self): file_to_attach = textwrap.dedent(""" import time @@ -249,7 +255,7 @@ def test_basic(self): os.remove("attached_script.py") -@unittest.skipIf(sys.platform == "darwin" and sys.version_info >= (3, 11), "Does not support 3.11+ on Mac") +@unittest.skipIf(attach_unavailable, "Does not support this platform") class TestAttachScript(CmdlineTmpl): def test_attach_script(self): # Isolate the attach stuff in a separate process