From 4fd473f55ef52d4222186be2b614bbd4d9c0fba0 Mon Sep 17 00:00:00 2001
From: Matt Wozniski <mwozniski@bloomberg.net>
Date: Wed, 21 Feb 2024 19:09:07 -0500
Subject: [PATCH] tests: Improve tests for our `sys.path` changes

As these were previously written, they didn't actually prove that the
original entry added by Python when memray was run had been removed,
which is especially problematic since:

    python -m memray run -m somemodule

would remove an entry added by Python and re-add that same entry. We had
no proof this was actually running the expected code and making the
expected changes.

Use `python -c ...` when running `memray run -m` to exercise this
better, since we can check that the (different) path that would be added
by the interpreter has been replaced.

Continue using `python -m memray` when running `memray run -c cmd` and
`memray run script.py`, but have the script that gets run print the
entirety of `sys.path` so we can make assertions about what isn't there
as well as what is.

Signed-off-by: Matt Wozniski <mwozniski@bloomberg.net>
---
 tests/integration/test_main.py | 57 +++++++++++++---------------------
 1 file changed, 22 insertions(+), 35 deletions(-)

diff --git a/tests/integration/test_main.py b/tests/integration/test_main.py
index ee9d3aa8fb..def46ef7df 100644
--- a/tests/integration/test_main.py
+++ b/tests/integration/test_main.py
@@ -268,9 +268,7 @@ def test_sys_manipulations_when_running_script(self, tmp_path):
         # GIVEN
         out_file = tmp_path / "result.bin"
         target_file = tmp_path / "test.py"
-        target_file.write_text("import some_adjacent_module")
-        other_file = tmp_path / "some_adjacent_module.py"
-        other_file.write_text("import sys; print(sys.argv); print(sys.path[0])")
+        target_file.write_text("import json, sys; print(json.dumps(sys.path))")
 
         # WHEN
         proc = subprocess.run(
@@ -292,57 +290,48 @@ def test_sys_manipulations_when_running_script(self, tmp_path):
         )
 
         # THEN
-        assert proc.returncode == 0
-        argv, path0 = proc.stdout.splitlines()
-        assert argv == repr([str(target_file), "some", "provided args"])
-        assert path0 == str(tmp_path)
         assert out_file.exists()
+        assert proc.returncode == 0
+        # Running `python -m` put cwd in sys.path; ensure we replaced it.
+        path = json.loads(proc.stdout)
+        assert os.getcwd() not in path
+        assert str(tmp_path) in path
 
     def test_sys_manipulations_when_running_module(self, tmp_path):
         # GIVEN
         out_file = tmp_path / "result.bin"
-        target_file = tmp_path / "test.py"
-        target_file.write_text("import some_adjacent_module")
-        other_file = tmp_path / "some_adjacent_module.py"
-        other_file.write_text("import sys; print(sys.argv); print(sys.path[0])")
-        env = os.environ.copy()
-        env["PYTHONPATH"] = str(tmp_path) + ":" + os.environ.get("PYTHONPATH", "")
 
         # WHEN
         proc = subprocess.run(
             [
                 sys.executable,
-                "-m",
-                "memray",
+                "-c",
+                "import sys; from memray.commands import main; sys.exit(main())",
                 "run",
                 "--quiet",
                 "--output",
                 str(out_file),
                 "-m",
-                "test",
-                "some",
-                "provided args",
+                "site",
             ],
             check=True,
             capture_output=True,
             text=True,
-            env=env,
         )
 
         # THEN
-        assert proc.returncode == 0
-        argv, path0 = proc.stdout.splitlines()
-        assert argv == repr([str(target_file), "some", "provided args"])
-        assert path0 == os.getcwd()
         assert out_file.exists()
+        assert proc.returncode == 0
+        # Running `python -c` put "" in sys.path; ensure we replaced it.
+        path = eval(
+            " ".join(line for line in proc.stdout.splitlines() if line.startswith(" "))
+        )
+        assert "" not in path
+        assert os.getcwd() in path
 
     def test_sys_manipulations_when_running_cmd(self, tmp_path):
         # GIVEN
         out_file = tmp_path / "result.bin"
-        other_file = tmp_path / "some_adjacent_module.py"
-        other_file.write_text("import sys; print(sys.argv); print(sys.path[0])")
-        env = os.environ.copy()
-        env["PYTHONPATH"] = str(tmp_path) + ":" + os.environ.get("PYTHONPATH", "")
 
         # WHEN
         proc = subprocess.run(
@@ -355,22 +344,20 @@ def test_sys_manipulations_when_running_cmd(self, tmp_path):
                 "--output",
                 str(out_file),
                 "-c",
-                "import some_adjacent_module",
-                "some",
-                "provided args",
+                "import json, sys; print(json.dumps(sys.path))",
             ],
             check=True,
             capture_output=True,
             text=True,
-            env=env,
         )
 
         # THEN
-        assert proc.returncode == 0
-        argv, path0 = proc.stdout.splitlines()
-        assert argv == repr(["-c", "some", "provided args"])
-        assert path0 == ""
         assert out_file.exists()
+        assert proc.returncode == 0
+        # Running `python -m` put cwd in sys.path; ensure we replaced it.
+        path = json.loads(proc.stdout)
+        assert os.getcwd() not in path
+        assert "" in path
 
     @pytest.mark.parametrize("option", [None, "--live", "--live-remote"])
     def test_run_file_that_is_not_python(self, capsys, option):