Migrate Managed Lustre to recording framework (#1337)

* Partial Managed Lustre migration

* Add helper method to RecordedCommandTestBase

* Complete work

* Remove ifdef trying to change polling, need a different design in the future

* Small change

* Small change

Author: alzimmermsft

core/Azure.Mcp.Core/tests/Azure.Mcp.Tests/Client/RecordedCommandTestsBase.cs

@@ -29,13 +29,13 @@ public abstract class RecordedCommandTestsBase(ITestOutputHelper output, TestPro
     /// <summary>
     /// Sanitizers that will apply generally across all parts (URI, Body, HeaderValues) of the request/response. This sanitization is applied to to recorded data at rest and during recording, and against test requests during playback.
     /// </summary>
-    public virtual List<GeneralRegexSanitizer> GeneralRegexSanitizers { get; } = new();
+    public virtual List<GeneralRegexSanitizer> GeneralRegexSanitizers { get; } = [];
 
     /// <summary>
     /// Sanitizers that will apply a regex to specific headers. This sanitization is applied to to recorded data at rest and during recording, and against test requests during playback.
     /// </summary>
-    public virtual List<HeaderRegexSanitizer> HeaderRegexSanitizers { get; } = new()
-    {
+    public virtual List<HeaderRegexSanitizer> HeaderRegexSanitizers { get; } =
+    [
         // Sanitize the WWW-Authenticate header which may contain tenant IDs or resource URLs to "Sanitized"
         // During conversion to recordings, the actual tenant ID is captured in group 1 and replaced with a fixed GUID.
         // REMOVAL of this formatting cause complete failure on tool side when it expects a valid URL with a GUID tenant ID.
@@ -46,37 +46,37 @@ public abstract class RecordedCommandTestsBase(ITestOutputHelper output, TestPro
             GroupForReplace = "1",
             Value = EmptyGuid
         })
-    };
+    ];
 
     /// <summary>
     /// Sanitizers that apply a regex replacement to URIs. This sanitization is applied to to recorded data at rest and during recording, and against test requests during playback.
     /// </summary>
-    public virtual List<UriRegexSanitizer> UriRegexSanitizers { get; } = new();
+    public virtual List<UriRegexSanitizer> UriRegexSanitizers { get; } = [];
 
     /// <summary>
     /// Sanitizers that will apply a regex replacement to a specific json body key. This sanitization is applied to to recorded data at rest and during recording, and against test requests during playback.
     /// </summary>
-    public virtual List<BodyKeySanitizer> BodyKeySanitizers { get; } = new();
+    public virtual List<BodyKeySanitizer> BodyKeySanitizers { get; } = [];
 
     /// <summary>
     /// Sanitizers that will apply regex replacement to the body of requests/responses. This sanitization is applied to to recorded data at rest and during recording, and against test requests during playback.
     /// </summary>
-    public virtual List<BodyRegexSanitizer> BodyRegexSanitizers { get; } = new();
+    public virtual List<BodyRegexSanitizer> BodyRegexSanitizers { get; } = [];
 
     /// <summary>
     /// The test-proxy has a default set of ~90 sanitizers for common sensitive data (GUIDs, tokens, timestamps, etc). This list allows opting out of specific default sanitizers by name.
     /// Grab the names from the test-proxy source at https://github.com/Azure/azure-sdk-tools/blob/main/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/SanitizerDictionary.cs#L65)
     /// Default Set:
     ///     - `AZSDK3430`: `$..id`
     /// </summary>
-    public virtual List<string> DisabledDefaultSanitizers { get; } = new() { "AZSDK3430" };
+    public virtual List<string> DisabledDefaultSanitizers { get; } = ["AZSDK3430"];
 
     /// <summary>
     /// During recording, variables saved to this dictionary will be propagated to the test-proxy and saved in the recording file.
     /// During playback, these variables will be available within the test function body, and can be used to ensure that dynamic values from the recording are used where
     /// specific values should be used.
     /// </summary>
-    protected readonly Dictionary<string, string> TestVariables = new Dictionary<string, string>();
+    protected readonly Dictionary<string, string> TestVariables = [];
 
     /// <summary>
     /// When set, applies a custom matcher for _all_ playback tests from this test class. This can be overridden on a per-test basis using the <see cref="CustomMatcherAttribute"/> attribute on test methods.
@@ -94,6 +94,31 @@ public virtual void RegisterVariable(string name, string value)
         TestVariables[name] = value;
     }
 
+    /// <summary>
+    /// Registers a variable to or retrieves it from the session record. This is a convenience equivalent to calling
+    /// <see cref="RegisterVariable(string, string)"/> and then retrieving the value from <see cref="TestVariables"/>.
+    /// If the test mode is playback, it will load attempt to load the variable and return it. If the test mode is
+    /// record, it will store the value and return it.
+    /// </summary>
+    /// <param name="name">The name of the variable to register or retrieve.</param>
+    /// <param name="value">The value reference to register or retrieve.</param>
+    /// <returns>The value of the variable.</returns>
+    public virtual string RegisterOrRetrieveVariable(string name, string value)
+    {
+        if (TestMode == TestMode.Record)
+        {
+            // store the value in the recording
+            TestVariables[name] = value;
+        }
+        else if (TestMode == TestMode.Playback)
+        {
+            // retrieve the value from the recording
+            value = TestVariables[name];
+        }
+
+        return value;
+    }
+
     // used to resolve a recording "path" given an invoking test
     protected static readonly RecordingPathResolver PathResolver = new();
 
@@ -238,7 +263,7 @@ private async Task DisableSanitizersAsync()
     {
         if (DisabledDefaultSanitizers.Count > 0)
         {
-            var toRemove = new SanitizerList(new List<string>());
+            var toRemove = new SanitizerList([]);
             foreach (var sanitizer in DisabledDefaultSanitizers)
             {
                 toRemove.Sanitizers.Add(sanitizer);

tools/Azure.Mcp.Tools.AppConfig/tests/Azure.Mcp.Tools.AppConfig.LiveTests/AppConfigCommandTests.cs

@@ -31,7 +31,7 @@ public class AppConfigCommandTests : RecordedCommandTestsBase
     /// AZSDK3493 = $..name
     /// AZSDK3447 = $..key
     /// </summary>
-    public override List<string> DisabledDefaultSanitizers => base.DisabledDefaultSanitizers.Concat(new[] { "AZSDK3493", "AZSDK3447" }).ToList();
+    public override List<string> DisabledDefaultSanitizers => [.. base.DisabledDefaultSanitizers, "AZSDK3493", "AZSDK3447"];
 
     public AppConfigCommandTests(ITestOutputHelper output, TestProxyFixture fixture) : base(output, fixture)
     {

tools/Azure.Mcp.Tools.KeyVault/tests/Azure.Mcp.Tools.KeyVault.LiveTests/KeyVaultCommandTests.cs

@@ -18,20 +18,20 @@ public class KeyVaultCommandTests(ITestOutputHelper output, TestProxyFixture fix
 {
     private readonly KeyVaultTestCertificateAssets _importCertificateAssets = KeyVaultTestCertificates.Load();
 
-    public override List<BodyRegexSanitizer> BodyRegexSanitizers => new List<BodyRegexSanitizer>() {
+    public override List<BodyRegexSanitizer> BodyRegexSanitizers => [
         // Sanitizes all hostnames in URLs to remove actual vault names (not limited to `kid` fields)
         new BodyRegexSanitizer(new BodyRegexSanitizerBody() {
           Regex = "(?<=http://|https://)(?<host>[^/?\\.]+)",
           GroupForReplace = "host",
         })
-    };
+    ];
 
     public override List<BodyKeySanitizer> BodyKeySanitizers
     {
         get
         {
-            return new List<BodyKeySanitizer>()
-            {
+            return
+            [
                 new BodyKeySanitizer(new BodyKeySanitizerBody("value")
                 {
                     Value = _importCertificateAssets.PfxBase64
@@ -44,7 +44,7 @@ public override List<BodyKeySanitizer> BodyKeySanitizers
                 {
                     Value = _importCertificateAssets.CsrBase64
                 })
-            };
+            ];
         }
     }
 
@@ -90,23 +90,21 @@ public async Task Should_get_key()
     [Fact]
     public async Task Should_create_key()
     {
-        var keyName = "key" + Random.Shared.NextInt64();
-
-        RegisterVariable("keyName", keyName);
+        var keyName = RegisterOrRetrieveVariable("keyName", "key" + Random.Shared.NextInt64());
 
         var result = await CallToolAsync(
             "keyvault_key_create",
             new()
             {
                 { "subscription", Settings.SubscriptionId },
                 { "vault", Settings.ResourceBaseName },
-                { "key", TestVariables["keyName"]},
+                { "key", keyName},
                 { "key-type", KeyType.Rsa.ToString() }
             });
 
         var createdKeyName = result.AssertProperty("name");
         Assert.Equal(JsonValueKind.String, createdKeyName.ValueKind);
-        Assert.Equal(TestVariables["keyName"], createdKeyName.GetString());
+        Assert.Equal(keyName, createdKeyName.GetString());
 
         var keyType = result.AssertProperty("keyType");
         Assert.Equal(JsonValueKind.String, keyType.ValueKind);
@@ -193,22 +191,20 @@ public async Task Should_get_certificate()
     [Fact]
     public async Task Should_create_certificate()
     {
-        var certificateName = "certificate" + Random.Shared.NextInt64();
-
-        RegisterVariable("certificateName", certificateName);
+        var certificateName = RegisterOrRetrieveVariable("certificateName", "certificate" + Random.Shared.NextInt64());
 
         var result = await CallToolAsync(
             "keyvault_certificate_create",
             new()
             {
                 { "subscription", Settings.SubscriptionId },
                 { "vault", Settings.ResourceBaseName },
-                { "certificate", TestVariables["certificateName"]}
+                { "certificate", certificateName}
             });
 
         var createdCertificateName = result.AssertProperty("name");
         Assert.Equal(JsonValueKind.String, createdCertificateName.ValueKind);
-        Assert.Equal(TestVariables["certificateName"], createdCertificateName.GetString());
+        Assert.Equal(certificateName, createdCertificateName.GetString());
 
         // Verify that the certificate has some expected properties
         ValidateCertificate(result);
@@ -224,23 +220,21 @@ public async Task Should_import_certificate()
 
         try
         {
-            var certificateName = "certificateimport" + Random.Shared.NextInt64();
-
-            RegisterVariable("certificateName", certificateName);
+            var certificateName = RegisterOrRetrieveVariable("certificateName", "certificateimport" + Random.Shared.NextInt64());
 
             var result = await CallToolAsync(
                 "keyvault_certificate_import",
                 new()
                 {
                     { "subscription", Settings.SubscriptionId },
                     { "vault", Settings.ResourceBaseName },
-                    { "certificate", TestVariables["certificateName"] },
+                    { "certificate", certificateName },
                     { "certificate-data", tempPath },
                     { "password", fakePassword }
                 });
             var createdCertificateName = result.AssertProperty("name");
             Assert.Equal(JsonValueKind.String, createdCertificateName.ValueKind);
-            Assert.Equal(TestVariables["certificateName"], createdCertificateName.GetString());
+            Assert.Equal(certificateName, createdCertificateName.GetString());
             // Validate basic certificate properties
             ValidateCertificate(result);
         }