[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));

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeFailoverTests/GlobalPartitionEndpointManagerTests.cs

@@ -385,8 +385,123 @@ public async Task TestRequestTimeoutExceptionScenarioAsync()
             Assert.AreEqual(HttpStatusCode.Created, response.StatusCode);
         }
 
+        [TestMethod]
+        [Timeout(20000)]
+        [Owner("dkunda")]
+        [DataRow(true, null, DisplayName = "Scenario when PPAF is enabled at client level and not set at the service level.")]
+        [DataRow(false, null, DisplayName = "Scenario when PPAF is disabled at client level and not set at the service level.")]
+        [DataRow(true, true, DisplayName = "Scenario when PPAF is enabled at client level and enabled at service level.")]
+        [DataRow(false, true, DisplayName = "Scenario when PPAF is disabled at client level and enabled at service level.")]
+        [DataRow(true, false, DisplayName = "Scenario when PPAF is enabled at client level and disabled at service level.")]
+        [DataRow(false, false, DisplayName = "Scenario when PPAF is disabled at client level and disabled at service level.")]
+        public async Task TestPPAFClientAndServerEnablementCombinationScenariosAsync(
+            bool ppafEnabledFromClient,
+            bool? ppafEnabledFromService)
+        {
+            if (ppafEnabledFromClient)
+            {
+                Environment.SetEnvironmentVariable(ConfigurationManager.PartitionLevelFailoverEnabled, "True");
+            }
+
+            try
+            {
+                GlobalPartitionEndpointManagerTests.SetupAccountAndCacheOperations(
+                    shouldEnablePPAF: ppafEnabledFromService,
+                    out string secondaryRegionNameForUri,
+                    out string globalEndpoint,
+                    out string secondaryRegionEndpiont,
+                    out string databaseName,
+                    out string containerName,
+                    out ResourceId containerResourceId,
+                    out Mock<IHttpHandler> mockHttpHandler,
+                    out IReadOnlyList<string> primaryRegionPartitionKeyRangeIds,
+                    out TransportAddressUri primaryRegionprimaryReplicaUri);
+
+                Mock<TransportClient> mockTransport = new Mock<TransportClient>(MockBehavior.Strict);
+
+                MockSetupsHelper.SetupServiceUnavailableException(
+                    mockTransport,
+                    primaryRegionprimaryReplicaUri);
+
+                // Partition key ranges are the same in both regions so the SDK
+                // does not need to go the secondary to get the partition key ranges.
+                // Only the addresses need to be mocked on the secondary
+                MockSetupsHelper.SetupAddresses(
+                    mockHttpHandler: mockHttpHandler,
+                    partitionKeyRangeId: primaryRegionPartitionKeyRangeIds.First(),
+                    regionEndpoint: secondaryRegionEndpiont,
+                    regionName: secondaryRegionNameForUri,
+                    containerResourceId: containerResourceId,
+                    primaryReplicaUri: out TransportAddressUri secondaryRegionPrimaryReplicaUri);
+
+                MockSetupsHelper.SetupCreateItemResponse(
+                    mockTransport,
+                    secondaryRegionPrimaryReplicaUri);
+
+                CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
+                {
+                    ConsistencyLevel = Cosmos.ConsistencyLevel.Strong,
+                    ApplicationPreferredRegions = new List<string>()
+                {
+                    Regions.EastUS,
+                    Regions.WestUS
+                },
+                    HttpClientFactory = () => new HttpClient(new HttpHandlerHelper(mockHttpHandler.Object)),
+                    TransportClientHandlerFactory = (original) => mockTransport.Object,
+                };
+
+                using CosmosClient customClient = new CosmosClient(
+                        globalEndpoint,
+                        Convert.ToBase64String(Encoding.UTF8.GetBytes(Guid.NewGuid().ToString())),
+                        cosmosClientOptions);
+
+                Container container = customClient.GetContainer(databaseName, containerName);
+
+                ToDoActivity toDoActivity = new ToDoActivity()
+                {
+                    Id = "TestItem",
+                    Pk = "TestPk"
+                };
+
+                if ((!ppafEnabledFromService.HasValue && ppafEnabledFromClient)
+                    || (ppafEnabledFromService.HasValue && ppafEnabledFromService.Value))
+                {
+                    ItemResponse<ToDoActivity> response = await container.CreateItemAsync(toDoActivity, new Cosmos.PartitionKey(toDoActivity.Pk));
+
+                    Assert.AreEqual(HttpStatusCode.Created, response.StatusCode);
+                    Assert.IsTrue(response.Diagnostics.GetContactedRegions().Count > 1);
+
+                    mockTransport.VerifyAll();
+                    mockHttpHandler.VerifyAll();
+                }
+                else
+                {
+                    try
+                    {
+                        await container.CreateItemAsync(toDoActivity, new Cosmos.PartitionKey(toDoActivity.Pk));
+                        Assert.Fail("Should throw an exception");
+                    }
+                    catch (CosmosException ce)
+                    {
+                        // Clears all the setups. No network calls should be done on the next operation.
+                        Assert.IsNotNull(ce);
+                        Assert.AreEqual(HttpStatusCode.ServiceUnavailable, ce.StatusCode);
+                    }
+                }
+
+                mockHttpHandler.Reset();
+                mockTransport.Reset();
+                mockTransport.Setup(x => x.Dispose());
+            }
+            finally
+            {
+                // Reset the environment variable to avoid affecting other tests.
+                Environment.SetEnvironmentVariable(ConfigurationManager.PartitionLevelFailoverEnabled, null);
+            }
+        }
+
         private static void SetupAccountAndCacheOperations(
-            bool shouldEnablePPAF,
+            bool? shouldEnablePPAF,
             out string secondaryRegionNameForUri,
             out string globalEndpoint,
             out string secondaryRegionEndpiont,

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeFailoverTests/MockSetupsHelper.cs

@@ -30,7 +30,7 @@ public static void SetupStrongAccountProperties(
             string endpoint,
             IList<AccountRegion> writeRegions,
             IList<AccountRegion> readRegions,
-            bool shouldEnablePPAF)
+            bool? shouldEnablePPAF)
         {
             HttpResponseMessage httpResponseMessage = MockSetupsHelper.CreateStrongAccount(
                 accountName,
@@ -109,7 +109,7 @@ public static HttpResponseMessage CreateStrongAccount(
             string accountName,
             IList<AccountRegion> writeRegions,
             IList<AccountRegion> readRegions,
-            bool shouldEnablePPAF = false)
+            bool? shouldEnablePPAF = null)
         {
             AccountProperties accountProperties = new AccountProperties()
             {