[Internal] PPAF: Fixes Integration Tests to Intercept Gateway Response

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemIntegrationTests.cs

@@ -5,14 +5,16 @@
     using System.IO;
     using System.Linq;
     using System.Net;
+    using System.Net.Http;
     using System.Text;
     using System.Text.Json;
     using System.Text.Json.Serialization;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Diagnostics;
     using Microsoft.Azure.Cosmos.FaultInjection;
-    using Microsoft.VisualStudio.TestTools.UnitTesting;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+    using Newtonsoft.Json.Linq;
     using static Microsoft.Azure.Cosmos.Routing.GlobalPartitionEndpointManagerCore;
     using static Microsoft.Azure.Cosmos.SDK.EmulatorTests.MultiRegionSetupHelpers;
 
@@ -1327,7 +1329,7 @@ public async Task CreateAndReadItemAsync_WithCircuitBreakerEnabledAndMultiMaster
 
         [TestMethod]
         [Owner("dkunda")]
-        [TestCategory("MultiRegion")]
+        [TestCategory("MultiRegion")]
         [Timeout(70000)]
         [DataRow(true, DisplayName = "Test scenario when PPAF is enabled at client level.")]
         [DataRow(false, DisplayName = "Test scenario when PPAF is disabled at client level.")]
@@ -1339,7 +1341,7 @@ public async Task ReadItemAsync_WithPPAFEnabledAndSingleMasterAccountWithRespons
             {
                 Environment.SetEnvironmentVariable(ConfigurationManager.PartitionLevelFailoverEnabled, "True");
             }
-
+
             // Enabling fault injection rule to simulate a 503 service unavailable scenario.
             string serviceUnavailableRuleId = "503-rule-" + Guid.NewGuid().ToString();
             FaultInjectionRule serviceUnavailableRule = new FaultInjectionRuleBuilder(
@@ -1356,15 +1358,38 @@ public async Task ReadItemAsync_WithPPAFEnabledAndSingleMasterAccountWithRespons
                 .Build();
 
             List<FaultInjectionRule> rules = new List<FaultInjectionRule> { serviceUnavailableRule };
-            FaultInjector faultInjector = new FaultInjector(rules);
+            FaultInjector faultInjector = new FaultInjector(rules);
+
+            // Now that the ppaf enablement flag is returned from gateway, we need to intercept the response and remove the flag from the response, so that
+            // the environment variable set above is honored.
+            HttpClientHandlerHelper httpClientHandlerHelper = new HttpClientHandlerHelper()
+            {
+                ResponseIntercepter = async (response, request) =>
+                {
+                    string json = await response?.Content?.ReadAsStringAsync();
+                    if (json.Length > 0 && json.Contains("enablePerPartitionFailoverBehavior"))
+                    {
+                        JObject parsedDatabaseAccountResponse = JObject.Parse(json);
+                        parsedDatabaseAccountResponse.Remove("enablePerPartitionFailoverBehavior");
+
+                        HttpResponseMessage interceptedResponse = Newtonsoft.Json.JsonConvert.DeserializeObject<HttpResponseMessage>(
+                            value: parsedDatabaseAccountResponse.ToString());
+
+                        return interceptedResponse;
+                    }
+
+                    return response;
+                },
+            };
 
             List<string> preferredRegions = new List<string> { region1, region2, region3 };
             CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
             {
                 ConsistencyLevel = ConsistencyLevel.Session,
                 FaultInjector = faultInjector,
                 RequestTimeout = TimeSpan.FromSeconds(5),
-                ApplicationPreferredRegions = preferredRegions,
+                ApplicationPreferredRegions = preferredRegions,
+                HttpClientFactory = () => new HttpClient(httpClientHandlerHelper),
             };
 
             List<CosmosIntegrationTestObject> itemsList = new()
@@ -1386,7 +1411,7 @@ public async Task ReadItemAsync_WithPPAFEnabledAndSingleMasterAccountWithRespons
 
                 ItemResponse<CosmosIntegrationTestObject> readResponse = await container.ReadItemAsync<CosmosIntegrationTestObject>(
                     id: itemsList[0].Id,
-                    partitionKey: new PartitionKey(itemsList[0].Pk));
+                    partitionKey: new PartitionKey(itemsList[0].Pk));
 
                 IReadOnlyList<(string regionName, Uri uri)> contactedRegionMapping = readResponse.Diagnostics.GetContactedRegions();
                 HashSet<string> contactedRegions = new(contactedRegionMapping.Select(r => r.regionName));