[Test Proxy] Add environment_variables fixture

sdk/conftest.py

@@ -26,7 +26,7 @@
 import os
 import pytest
 
-from devtools_testutils import recorded_test, test_proxy, variable_recorder
+from devtools_testutils import environment_variables, recorded_test, test_proxy, variable_recorder
 
 
 def pytest_configure(config):

tools/azure-sdk-tools/devtools_testutils/__init__.py

@@ -18,7 +18,7 @@
 # cSpell:disable
 from .envvariable_loader import EnvironmentVariableLoader
 PowerShellPreparer = EnvironmentVariableLoader  # Backward compat
-from .proxy_fixtures import recorded_test, variable_recorder
+from .proxy_fixtures import environment_variables, recorded_test, variable_recorder
 from .proxy_startup import start_test_proxy, stop_test_proxy, test_proxy
 from .proxy_testcase import recorded_by_proxy
 from .sanitizers import (
@@ -88,6 +88,7 @@
     "PemCertificate",
     "PowerShellPreparer",
     "EnvironmentVariableLoader",
+    "environment_variables",
     "recorded_by_proxy",
     "recorded_test",
     "test_proxy",

tools/azure-sdk-tools/devtools_testutils/proxy_fixtures.py

@@ -5,6 +5,7 @@
 # --------------------------------------------------------------------------
 from inspect import iscoroutinefunction
 import logging
+import os
 from typing import TYPE_CHECKING
 import urllib.parse as url_parse
 
@@ -16,15 +17,102 @@
 # the functions we patch
 from azure.core.pipeline.transport import RequestsTransport
 
-from .helpers import get_test_id, is_live_and_not_recording
+from .helpers import get_test_id, is_live, is_live_and_not_recording
 from .proxy_testcase import start_record_or_playback, stop_record_or_playback, transform_request
 from .proxy_startup import test_proxy
+from .sanitizers import add_general_regex_sanitizer
 
 if TYPE_CHECKING:
     from typing import Any, Callable, Dict, Optional, Tuple
     from pytest import FixtureRequest
 
 
+_LOGGER = logging.getLogger()
+
+
+class EnvironmentVariableSanitizer:
+    def __init__(self) -> None:
+        self._fake_values = {}
+
+    def sanitize(self, variable: str, value: str) -> str:
+        """Registers a sanitizer that replaces the value of the specified environment variable with the provided value.
+
+        :param str variable: Name of the environment variable to sanitize.
+        :param str value: Value to sanitize the environment variable's value with.
+
+        :returns: The real value of `variable` in live mode, or the sanitized value in playback.
+        """
+        self._fake_values[variable] = value
+        real_value = os.getenv(variable)
+        if real_value:
+            add_general_regex_sanitizer(regex=real_value, value=value)
+        else:
+            _LOGGER.info(f"No value for {variable} was found, so a sanitizer could not be registered for the variable.")
+
+        return real_value if is_live() else value
+
+    def sanitize_batch(self, variables: "Dict[str, str]") -> "Dict[str, str]":
+        """Registers sanitizers that replace the values of multiple envrionment variables with the provided values.
+
+        :param variables: A dictionary mapping environment variable names to values they should be sanitized with.
+            For example: {"SERICE_CLIENT_ID": "fake_client_id", "SERVICE_ENDPOINT": "https://fake-endpoint.azure.net"}
+
+        :returns: A dictionary mapping environment variables to their real values in live mode, or their sanitized
+            values in playback.
+        """
+        current_values = {variable: self.sanitize(variable, variables[variable]) for variable in variables}
+        return current_values
+
+
+    def get(self, variable: str) -> str:
+        """Returns the value of the specified environment variable in live mode, or the sanitized value in playback.
+
+        :param str variable: Name of the environment variable to fetch the value of.
+
+        :returns: The real value of `variable` in live mode, or the sanitized value in playback.
+        """
+        return os.getenv(variable) if is_live() else self._fake_values.get(variable)
+
+
+class VariableRecorder:
+    def __init__(self, variables: "Dict[str, str]") -> None:
+        self.variables = variables
+
+    def get_or_record(self, variable: str, default: str) -> str:
+        """Returns the recorded value of `variable`, or records and returns `default` as the value for `variable`.
+
+        In recording mode, `get_or_record("a", "b")` will record "b" for the value of the variable `a` and return "b".
+        In playback, it will return the recorded value of `a`. This is an analogue of a Python dictionary's `setdefault`
+        method: https://docs.python.org/library/stdtypes.html#dict.setdefault.
+
+        :param str variable: The name of the variable to search the value of, or record a value for.
+        :param str default: The variable value to record.
+
+        :returns: str
+        """
+        if not isinstance(default, str):
+            raise ValueError('"default" must be a string. The test proxy cannot record non-string variable values.')
+        return self.variables.setdefault(variable, default)
+
+
+@pytest.fixture
+def environment_variables(test_proxy: None) -> EnvironmentVariableSanitizer:
+    """Fixture that returns an EnvironmentVariableSanitizer for convenient environment variable fetching and sanitizing.
+
+    :param test_proxy: The fixture responsible for starting up the test proxy server.
+    :type test_proxy: None
+
+    :returns: An EnvironmentVariableSanitizer object. Calling:
+        - `sanitize(a, b)` will sanitize the value of environnment variable `a` with value `b`
+        - `sanitize_batch(dict)` will sanitize the values of all variables in dictionary `dict`
+        - `get(a)` will return the value of environment variable `a` in the current context (live or playback mode)
+        See the definition of EnvironmentVariableSanitizer in
+        https://github.com/Azure/azure-sdk-for-python/blob/main/tools/azure-sdk-tools/devtools_testutils/proxy_fixtures.py
+        for more details.
+    """
+    return EnvironmentVariableSanitizer()
+
+
 @pytest.fixture
 async def recorded_test(test_proxy: None, request: "FixtureRequest") -> "Dict[str, Any]":
     """Fixture that redirects network requests to target the azure-sdk-tools test proxy.
@@ -56,44 +144,25 @@ async def recorded_test(test_proxy: None, request: "FixtureRequest") -> "Dict[st
 
 
 @pytest.fixture
-def variable_recorder(recorded_test: "Dict[str, Any]") -> "Dict[str, str]":
+def variable_recorder(recorded_test: "Dict[str, Any]") -> VariableRecorder:
     """Fixture that invokes the `recorded_test` fixture and returns a dictionary of recorded test variables.
 
     :param recorded_test: The fixture responsible for redirecting network traffic to target the test proxy.
         This should return a dictionary containing information about the current test -- in particular, the variables
         that were recorded with the test.
     :type recorded_test: Dict[str, Any]
 
-    :returns: A dictionary that maps test variables to string values. If no variable dictionary was stored when the test
-        was recorded, this returns an empty dictionary.
+    :returns: A VariableRecorder object. Calling `get_or_record(a, b)` on this object will return the recorded value of
+        `a` in playback mode, or record the value `b` in recording mode. See the definition of VariableRecorder in
+        https://github.com/Azure/azure-sdk-for-python/blob/main/tools/azure-sdk-tools/devtools_testutils/proxy_fixtures.py
+        for more details.
     """
     return VariableRecorder(recorded_test["variables"])
 
 
 # ----------HELPERS----------
 
 
-class VariableRecorder():
-    def __init__(self, variables: "Dict[str, str]") -> None:
-        self.variables = variables
-
-    def get_or_record(self, variable: str, default: str) -> str:
-        """Returns the recorded value of `variable`, or records and returns `default` as the value for `variable`.
-
-        In recording mode, `get_or_record("a", "b")` will record "b" for the value of the variable `a` and return "b".
-        In playback, it will return the recorded value of `a`. This is an analogue of a Python dictionary's `setdefault`
-        method: https://docs.python.org/library/stdtypes.html#dict.setdefault.
-
-        :param str variable: The name of the variable to search the value of, or record a value for.
-        :param str default: The variable value to record.
-
-        :returns: str
-        """
-        if not isinstance(default, str):
-            raise ValueError('"default" must be a string. The test proxy cannot record non-string variable values.')
-        return self.variables.setdefault(variable, default)
-
-
 def start_proxy_session() -> "Optional[Tuple[str, str, Dict[str, str]]]":
     """Begins a playback or recording session and returns the current test ID, recording ID, and recorded variables.
 
@@ -205,8 +274,7 @@ def restore_async_traffic(original_transport_func: "Callable", request: "Fixture
         if isinstance(error, ResourceNotFoundError):
             error_body = ContentDecodePolicy.deserialize_from_http_generics(error.response)
             message = error_body.get("message") or error_body.get("Message")
-            logger = logging.getLogger()
-            logger.error(f"\n\n-----Test proxy playback error:-----\n\n{message}")
+            _LOGGER.error(f"\n\n-----Test proxy playback error:-----\n\n{message}")
 
 
 def restore_traffic(original_transport_func: "Callable", request: "FixtureRequest") -> None:
@@ -229,5 +297,4 @@ def restore_traffic(original_transport_func: "Callable", request: "FixtureReques
         if isinstance(error, ResourceNotFoundError):
             error_body = ContentDecodePolicy.deserialize_from_http_generics(error.response)
             message = error_body.get("message") or error_body.get("Message")
-            logger = logging.getLogger()
-            logger.error(f"\n\n-----Test proxy playback error:-----\n\n{message}")
+            _LOGGER.error(f"\n\n-----Test proxy playback error:-----\n\n{message}")