From bb3a48cdff5fc6430ae1fdcdf8d673606cd07076 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?McCoy=20Pati=C3=B1o?= <mcpatino@microsoft.com>
Date: Wed, 20 Jul 2022 09:07:31 -0700
Subject: [PATCH] Variable recording class

---
 .../devtools_testutils/__init__.py            |  3 +-
 .../devtools_testutils/proxy_testcase.py      | 15 ------
 .../devtools_testutils/variable_recorder.py   | 49 +++++++++++++++++++
 3 files changed, 51 insertions(+), 16 deletions(-)
 create mode 100644 tools/azure-sdk-tools/devtools_testutils/variable_recorder.py

diff --git a/tools/azure-sdk-tools/devtools_testutils/__init__.py b/tools/azure-sdk-tools/devtools_testutils/__init__.py
index 8aad9bad74dd..9d57cd551d20 100644
--- a/tools/azure-sdk-tools/devtools_testutils/__init__.py
+++ b/tools/azure-sdk-tools/devtools_testutils/__init__.py
@@ -19,7 +19,7 @@
 from .envvariable_loader import EnvironmentVariableLoader
 PowerShellPreparer = EnvironmentVariableLoader  # Backward compat
 from .proxy_startup import start_test_proxy, stop_test_proxy, test_proxy
-from .proxy_testcase import recorded_by_proxy, recorded_test, variable_recorder
+from .proxy_testcase import recorded_by_proxy, recorded_test
 from .sanitizers import (
     add_body_key_sanitizer,
     add_body_regex_sanitizer,
@@ -34,6 +34,7 @@
     set_custom_default_matcher,
     set_default_settings,
 )
+from .variable_recorder import variable_recorder
 from .helpers import ResponseCallback, RetryCounter
 from .fake_credentials import FakeTokenCredential
 
diff --git a/tools/azure-sdk-tools/devtools_testutils/proxy_testcase.py b/tools/azure-sdk-tools/devtools_testutils/proxy_testcase.py
index 691c5551b0d3..cf1cfc7ec3de 100644
--- a/tools/azure-sdk-tools/devtools_testutils/proxy_testcase.py
+++ b/tools/azure-sdk-tools/devtools_testutils/proxy_testcase.py
@@ -365,18 +365,3 @@ def restore_traffic(original_transport_func: "Callable", request: "FixtureReques
             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}")
-
-
-@pytest.fixture
-def variable_recorder(recorded_test: "Dict[str, Any]") -> "Dict[str, str]":
-    """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.
-    """
-    return recorded_test["variables"]
diff --git a/tools/azure-sdk-tools/devtools_testutils/variable_recorder.py b/tools/azure-sdk-tools/devtools_testutils/variable_recorder.py
new file mode 100644
index 000000000000..ba5b5032cacf
--- /dev/null
+++ b/tools/azure-sdk-tools/devtools_testutils/variable_recorder.py
@@ -0,0 +1,49 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+from typing import TYPE_CHECKING
+
+import pytest
+
+from .proxy_testcase import recorded_test
+
+if TYPE_CHECKING:
+    from typing import Any, Dict
+
+
+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 variable_recorder(recorded_test: "Dict[str, Any]") -> "Dict[str, str]":
+    """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.
+    """
+    return recorded_test["variables"]