Save env vars at first pytest hook to ensure theyre unedited (#22344)

fixes #22192. Now, all environment variables are accessed during the pytest_load_initial_conftests hook and then saved as global variables in the plugin. This ensures the port and uuid will still be saved even if a user safely or unsafely clears their environment variables during testing. --------- Co-authored-by: Karthik Nadig <[email protected]>
microsoft · Oct 26, 2023 · 71a451c · 71a451c
1 parent c2dec14
commit 71a451c
Show file tree

Hide file tree

Showing 4 changed files with 146 additions and 40 deletions.
diff --git a/pythonFiles/tests/pytestadapter/.data/test_env_vars.py b/pythonFiles/tests/pytestadapter/.data/test_env_vars.py
@@ -0,0 +1,32 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+
+
+def test_clear_env(monkeypatch):
+    # Clear all environment variables
+    monkeypatch.setattr(os, "environ", {})
+
+    # Now os.environ should be empty
+    assert not os.environ
+
+    # After the test finishes, the environment variables will be reset to their original state
+
+
+def test_check_env():
+    # This test will have access to the original environment variables
+    assert "PATH" in os.environ
+
+
+def test_clear_env_unsafe():
+    # Clear all environment variables
+    os.environ.clear()
+    # Now os.environ should be empty
+    assert not os.environ
+
+
+def test_check_env_unsafe():
+    # ("PATH" in os.environ) is False here if it runs after test_clear_env_unsafe.
+    # Regardless, this test will pass and TEST_PORT and TEST_UUID will still be set correctly
+    assert "PATH" not in os.environ
diff --git a/pythonFiles/tests/pytestadapter/expected_execution_test_output.py b/pythonFiles/tests/pytestadapter/expected_execution_test_output.py
@@ -624,3 +624,63 @@
         "subtest": None,
     },
 }
+
+# This is the expected output for the test safe clear env vars file.
+# └── test_env_vars.py
+#    └── test_clear_env: success
+#    └── test_check_env: success
+
+test_safe_clear_env_vars_path = TEST_DATA_PATH / "test_env_vars.py"
+safe_clear_env_vars_expected_execution_output = {
+    get_absolute_test_id(
+        "test_env_vars.py::test_clear_env", test_safe_clear_env_vars_path
+    ): {
+        "test": get_absolute_test_id(
+            "test_env_vars.py::test_clear_env", test_safe_clear_env_vars_path
+        ),
+        "outcome": "success",
+        "message": None,
+        "traceback": None,
+        "subtest": None,
+    },
+    get_absolute_test_id(
+        "test_env_vars.py::test_check_env", test_safe_clear_env_vars_path
+    ): {
+        "test": get_absolute_test_id(
+            "test_env_vars.py::test_check_env", test_safe_clear_env_vars_path
+        ),
+        "outcome": "success",
+        "message": None,
+        "traceback": None,
+        "subtest": None,
+    },
+}
+
+# This is the expected output for the test unsafe clear env vars file.
+# └── test_env_vars.py
+#    └── test_clear_env_unsafe: success
+#    └── test_check_env_unsafe: success
+unsafe_clear_env_vars_expected_execution_output = {
+    get_absolute_test_id(
+        "test_env_vars.py::test_clear_env_unsafe", test_safe_clear_env_vars_path
+    ): {
+        "test": get_absolute_test_id(
+            "test_env_vars.py::test_clear_env_unsafe", test_safe_clear_env_vars_path
+        ),
+        "outcome": "success",
+        "message": None,
+        "traceback": None,
+        "subtest": None,
+    },
+    get_absolute_test_id(
+        "test_env_vars.py::test_check_env_unsafe", test_safe_clear_env_vars_path
+    ): {
+        "test": get_absolute_test_id(
+            "test_env_vars.py::test_check_env_unsafe", test_safe_clear_env_vars_path
+        ),
+        "outcome": "success",
+        "message": None,
+        "traceback": None,
+        "subtest": None,
+    },
+}
diff --git a/pythonFiles/tests/pytestadapter/test_execution.py b/pythonFiles/tests/pytestadapter/test_execution.py
@@ -131,6 +131,13 @@ def test_bad_id_error_execution():
 @pytest.mark.parametrize(
     "test_ids, expected_const",
     [
+        (
+            [
+                "test_env_vars.py::test_clear_env",
+                "test_env_vars.py::test_check_env",
+            ],
+            expected_execution_test_output.safe_clear_env_vars_expected_execution_output,
+        ),
         (
             [
                 "skip_tests.py::test_something",

diff --git a/pythonFiles/vscode_pytest/__init__.py b/pythonFiles/vscode_pytest/__init__.py
@@ -6,7 +6,6 @@
 import os
 import pathlib
 import sys
-import time
 import traceback
 
 import pytest
@@ -56,9 +55,21 @@ def __init__(self, message):
 IS_DISCOVERY = False
 map_id_to_path = dict()
 collected_tests_so_far = list()
+TEST_PORT = os.getenv("TEST_PORT")
+TEST_UUID = os.getenv("TEST_UUID")
 
 
 def pytest_load_initial_conftests(early_config, parser, args):
+    global TEST_PORT
+    global TEST_UUID
+    TEST_PORT = os.getenv("TEST_PORT")
+    TEST_UUID = os.getenv("TEST_UUID")
+    error_string = (
+        "PYTEST ERROR: TEST_UUID and/or TEST_PORT are not set at the time of pytest starting. Please confirm these environment variables are not being"
+        " changed or removed as they are required for successful test discovery and execution."
+        f" \nTEST_UUID = {TEST_UUID}\nTEST_PORT = {TEST_PORT}\n"
+    )
+    print(error_string, file=sys.stderr)
     if "--collect-only" in args:
         global IS_DISCOVERY
         IS_DISCOVERY = True
@@ -689,58 +700,54 @@ def send_post_request(
     payload -- the payload data to be sent.
     cls_encoder -- a custom encoder if needed.
     """
-    testPort = os.getenv("TEST_PORT")
-    testUuid = os.getenv("TEST_UUID")
-    if testPort is None:
-        print(
-            "Error[vscode-pytest]: TEST_PORT is not set.",
-            " TEST_UUID = ",
-            testUuid,
-        )
-        testPort = DEFAULT_PORT
-    if testUuid is None:
-        print(
-            "Error[vscode-pytest]: TEST_UUID is not set.",
-            " TEST_PORT = ",
-            testPort,
+    global TEST_PORT
+    global TEST_UUID
+    if TEST_UUID is None or TEST_PORT is None:
+        # if TEST_UUID or TEST_PORT is None, print an error and fail as these are both critical errors
+        error_msg = (
+            "PYTEST ERROR: TEST_UUID and/or TEST_PORT are not set at the time of pytest starting. Please confirm these environment variables are not being"
+            " changed or removed as they are required for successful pytest discovery and execution."
+            f" \nTEST_UUID = {TEST_UUID}\nTEST_PORT = {TEST_PORT}\n"
         )
-        testUuid = "unknown"
-    addr = ("localhost", int(testPort))
+        print(error_msg, file=sys.stderr)
+        raise VSCodePytestError(error_msg)
+
+    addr = ("localhost", int(TEST_PORT))
     global __socket
 
     if __socket is None:
         try:
             __socket = socket_manager.SocketManager(addr)
             __socket.connect()
         except Exception as error:
-            print(f"Plugin error connection error[vscode-pytest]: {error}")
+            error_msg = f"Error attempting to connect to extension communication socket[vscode-pytest]: {error}"
+            print(error_msg, file=sys.stderr)
+            print(
+                "If you are on a Windows machine, this error may be occurring if any of your tests clear environment variables"
+                " as they are required to communicate with the extension. Please reference https://docs.pytest.org/en/stable/how-to/monkeypatch.html#monkeypatching-environment-variables"
+                "for the correct way to clear environment variables during testing.\n",
+                file=sys.stderr,
+            )
             __socket = None
+            raise VSCodePytestError(error_msg)
 
     data = json.dumps(payload, cls=cls_encoder)
     request = f"""Content-Length: {len(data)}
 Content-Type: application/json
-Request-uuid: {testUuid}
+Request-uuid: {TEST_UUID}
 
 {data}"""
 
-    max_retries = 3
-    retries = 0
-    while retries < max_retries:
-        try:
-            if __socket is not None and __socket.socket is not None:
-                __socket.socket.sendall(request.encode("utf-8"))
-                # print("Post request sent successfully!")
-                # print("data sent", payload, "end of data")
-                break  # Exit the loop if the send was successful
-            else:
-                print("Plugin error connection error[vscode-pytest]")
-                print(f"[vscode-pytest] data: {request}")
-        except Exception as error:
-            print(f"Plugin error connection error[vscode-pytest]: {error}")
-            print(f"[vscode-pytest] data: {request}")
-            retries += 1  # Increment retry counter
-            if retries < max_retries:
-                print(f"Retrying ({retries}/{max_retries}) in 2 seconds...")
-                time.sleep(2)  # Wait for a short duration before retrying
-            else:
-                print("Maximum retry attempts reached. Cannot send post request.")
+    try:
+        if __socket is not None and __socket.socket is not None:
+            __socket.socket.sendall(request.encode("utf-8"))
+        else:
+            print(
+                f"Plugin error connection error[vscode-pytest], socket is None \n[vscode-pytest] data: \n{request} \n",
+                file=sys.stderr,
+            )
+    except Exception as error:
+        print(
+            f"Plugin error, exception thrown while attempting to send data[vscode-pytest]: {error} \n[vscode-pytest] data: \n{request}\n",
+            file=sys.stderr,
+        )