Modified Batch Track 2 Test format and included sanitizers with custo…

…m matchers (#35470) * Extended BatchTestBase from TestProxyTestBase and added sanitzations * Added sanitization rule for http url * Updated azure-core-test versions and migrated test files outside of generated folder * Added custom playback rules and matchers * Added test recording json files * Refactored Shared Key test setup
Azure · Apr 26, 2024 · c12bfc4 · c12bfc4
1 parent 33649fd
commit c12bfc4
Show file tree

Hide file tree

Showing 14 changed files with 5,320 additions and 47 deletions.
diff --git a/sdk/batch/azure-compute-batch/pom.xml b/sdk/batch/azure-compute-batch/pom.xml
@@ -48,7 +48,7 @@
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-core</artifactId>
-      <version>1.37.0</version> <!-- {x-version-update;com.azure:azure-core;dependency} -->
+      <version>1.40.0</version> <!-- {x-version-update;com.azure:azure-core;dependency} -->
     </dependency>
     <dependency>
       <groupId>com.azure</groupId>
@@ -76,7 +76,7 @@
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-core-test</artifactId>
-      <version>1.15.0</version> <!-- {x-version-update;com.azure:azure-core-test;dependency} -->
+      <version>1.18.0</version> <!-- {x-version-update;com.azure:azure-core-test;dependency} -->
       <scope>test</scope>
     </dependency>
     <dependency>

diff --git a/...generated/BatchServiceClientTestBase.java → ...ute/batch/BatchServiceClientTestBase.java b/...generated/BatchServiceClientTestBase.java → ...ute/batch/BatchServiceClientTestBase.java
@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 // Code generated by Microsoft (R) AutoRest Code Generator.
 
-package com.azure.compute.batch.generated;
+package com.azure.compute.batch;
 
 import com.azure.compute.batch.*;
 import com.azure.compute.batch.auth.BatchSharedKeyCredentials;
@@ -12,8 +12,13 @@
 import com.azure.core.http.policy.HttpLogDetailLevel;
 import com.azure.core.http.policy.HttpLogOptions;
 import com.azure.core.http.rest.PagedIterable;
-import com.azure.core.test.TestBase;
+import com.azure.core.test.InterceptorManager;
 import com.azure.core.test.TestMode;
+import com.azure.core.test.TestProxyTestBase;
+import com.azure.core.test.models.CustomMatcher;
+import com.azure.core.test.models.TestProxyRequestMatcher;
+import com.azure.core.test.models.TestProxySanitizer;
+import com.azure.core.test.models.TestProxySanitizerType;
 import com.azure.core.util.Configuration;
 import com.azure.identity.DefaultAzureCredentialBuilder;
 
@@ -42,7 +47,7 @@
 import org.junit.Assert;
 
 
-class BatchServiceClientTestBase extends TestBase {
+class BatchServiceClientTestBase extends TestProxyTestBase {
     protected BatchServiceClientBuilder batchClientBuilder;
 
     protected ApplicationsClient applicationsClient;
@@ -67,12 +72,15 @@ class BatchServiceClientTestBase extends TestBase {
 
 	static final int MAX_LEN_ID = 64;
 
+    static String REDACTED = "REDACTED";
+
 	public enum AuthMode {
 	        AAD, SharedKey
 	    }
-	
+
     @Override
     protected void beforeTest() {
+        super.beforeTest();
     	batchClientBuilder =
                 new BatchServiceClientBuilder()
                         .endpoint(Configuration.getGlobalConfiguration().get("AZURE_BATCH_ENDPOINT", "endpoint"))
@@ -82,8 +90,11 @@ protected void beforeTest() {
         	batchClientBuilder
                     .httpClient(interceptorManager.getPlaybackClient())
                     .credential(request -> Mono.just(new AccessToken("this_is_a_token", OffsetDateTime.MAX)));
+
+            addTestRulesOnPlayback(interceptorManager);
         } else if (getTestMode() == TestMode.RECORD) {
         	batchClientBuilder.addPolicy(interceptorManager.getRecordPolicy());
+            addTestSanitizersAndRules(interceptorManager);
         }
 
         authenticateClient(AuthMode.AAD);
@@ -99,7 +110,21 @@ protected void beforeTest() {
         computeNodesClient = batchClientBuilder.buildComputeNodesClient();
         computeNodeExtensionsClient = batchClientBuilder.buildComputeNodeExtensionsClient();
     }
-
+
+    public void addTestRulesOnPlayback(InterceptorManager interceptorManager) {
+        List<TestProxyRequestMatcher> customMatchers = new ArrayList<>();
+        customMatchers.add(new CustomMatcher().setComparingBodies(false));
+        customMatchers.add(new CustomMatcher().setExcludedHeaders(Collections.singletonList("ocp-date")));
+        interceptorManager.addMatchers(customMatchers);
+    }
+
+    public void addTestSanitizersAndRules(InterceptorManager interceptorManager) {
+        List<TestProxySanitizer> testProxySanitizers = new ArrayList<TestProxySanitizer>();
+        testProxySanitizers.add(new TestProxySanitizer("$..httpUrl", null,REDACTED, TestProxySanitizerType.BODY_KEY));
+        testProxySanitizers.add(new TestProxySanitizer("$..containerUrl", null,REDACTED, TestProxySanitizerType.BODY_KEY));
+        interceptorManager.addSanitizers(testProxySanitizers);
+    }
+
     void authenticateClient(AuthMode auth) {
 	    if (getTestMode() == TestMode.RECORD) {
             if (auth == AuthMode.AAD) {
@@ -138,7 +163,7 @@ static String getStringIdWithUserNamePrefix(String name) {
         }
         return out.toString();
     }
-    
+
     BatchPool createIfNotExistIaaSPool(String poolId) throws Exception {
         // Create a pool with 3 Small VMs
         String poolVmSize = "STANDARD_D1_V2";
@@ -147,7 +172,7 @@ BatchPool createIfNotExistIaaSPool(String poolId) throws Exception {
         // 10 minutes
         long poolSteadyTimeoutInSeconds = 10 * 60 * 1000;
         PoolClient poolClient = batchClientBuilder.buildPoolClient();
-        
+
 
         // Check if pool exists
         if (!poolExists(poolClient, poolId)) {
@@ -164,21 +189,21 @@ BatchPool createIfNotExistIaaSPool(String poolId) throws Exception {
 
          // Need VNet to allow security to inject NSGs
             NetworkConfiguration networkConfiguration = createNetworkConfiguration();
-            
+
             BatchPool poolToAdd = new BatchPool();
             poolToAdd.setId(poolId).setTargetDedicatedNodes(poolVmCount).setVmSize(poolVmSize)
             		 .setVirtualMachineConfiguration(configuration)
             		 .setUserAccounts(userList)
             		 .setNetworkConfiguration(networkConfiguration);
 
             poolClient.add(poolToAdd);
-        } 
+        }
         else {
         	System.out.println(String.format("The %s already exists.", poolId));
             //logger.log(createLogRecord(Level.INFO, String.format("The %s already exists.", poolId)));
         }
-        
-        
+
+
         long startTime = System.currentTimeMillis();
         long elapsedTime = 0L;
         boolean steady = false;
@@ -200,7 +225,7 @@ BatchPool createIfNotExistIaaSPool(String poolId) throws Exception {
 
         return pool;
     }
-    
+
    NetworkConfiguration createNetworkConfiguration(){
     	Configuration localConfig = Configuration.getGlobalConfiguration();
         String vnetName = localConfig.get("AZURE_VNET", "");
@@ -219,12 +244,12 @@ NetworkConfiguration createNetworkConfiguration(){
 
         	PagedIterable<Network> networks = manager.networks().listByResourceGroup(vnetResourceGroup);
         	boolean networksFound = false;
-        	
+
         	for (Network network: networks) {
         		networksFound = true;
         		break;
         	}
-        	
+
         	if (!networksFound) {
         		Network network = manager.networks().define(vnetName)
         				.withRegion(localConfig.get("AZURE_BATCH_REGION"))
@@ -244,15 +269,15 @@ NetworkConfiguration createNetworkConfiguration(){
 
         return new NetworkConfiguration().setSubnetId(vNetResourceId);
     }
-    
+
     void threadSleepInRecordMode(long millis) throws InterruptedException {
         // Called for long timeouts which should only happen in Record mode.
         // Speeds up the tests in Playback mode.
         if (getTestMode() == TestMode.RECORD) {
             Thread.sleep(millis);
         }
     }
-    
+
     static boolean poolExists(PoolClient poolClient, String poolId) {
     	try {
             poolClient.exists(poolId);

diff --git a/...ure/compute/batch/generated/JobTests.java → ...ava/com/azure/compute/batch/JobTests.java b/...ure/compute/batch/generated/JobTests.java → ...ava/com/azure/compute/batch/JobTests.java
@@ -1,12 +1,11 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
-package com.azure.compute.batch.generated;
+package com.azure.compute.batch;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 
-import com.azure.compute.batch.JobClient;
 import com.azure.compute.batch.models.AutoPoolSpecification;
 import com.azure.compute.batch.models.BatchJob;
 import com.azure.compute.batch.models.BatchPool;

diff --git a/...re/compute/batch/generated/PoolTests.java → ...va/com/azure/compute/batch/PoolTests.java b/...re/compute/batch/generated/PoolTests.java → ...va/com/azure/compute/batch/PoolTests.java
@@ -1,9 +1,8 @@
-package com.azure.compute.batch.generated;
+package com.azure.compute.batch;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
-import com.azure.compute.batch.ComputeNodesClient;
 import com.azure.compute.batch.models.AllocationState;
 import com.azure.compute.batch.models.BatchPool;
 import com.azure.compute.batch.models.BatchPoolResizeParameters;

diff --git a/...mpute/batch/generated/SharedKeyTests.java → ...m/azure/compute/batch/SharedKeyTests.java b/...mpute/batch/generated/SharedKeyTests.java → ...m/azure/compute/batch/SharedKeyTests.java
@@ -1,7 +1,5 @@
-package com.azure.compute.batch.generated;
+package com.azure.compute.batch;
 
-import com.azure.compute.batch.BatchServiceClientBuilder;
-import com.azure.compute.batch.PoolClient;
 import com.azure.compute.batch.auth.BatchSharedKeyCredentials;
 import com.azure.compute.batch.models.*;
 import com.azure.core.credential.AccessToken;
@@ -27,31 +25,17 @@
 import static java.time.OffsetDateTime.now;
 
 public class SharedKeyTests extends BatchServiceClientTestBase {
-    private static BatchServiceClientBuilder clientBuilderWithSharedKey;
     private static PoolClient poolClientWithSharedKey;
     private final String sharedKeyPoolId = "SharedKey-testpool";
     private final String vmSize = "STANDARD_D1_V2";
     private final String nodeAgentSkuId = "batch.node.ubuntu 18.04";
 
     @Override
     protected void beforeTest() {
+        super.beforeTest();
         BatchSharedKeyCredentials sharedKeyCred = getSharedKeyCredentials();
-        clientBuilderWithSharedKey =
-                new BatchServiceClientBuilder()
-                        .endpoint(Configuration.getGlobalConfiguration().get("AZURE_BATCH_ENDPOINT", "endpoint"))
-                        .httpClient(HttpClient.createDefault())
-                        .httpLogOptions(new HttpLogOptions().setLogLevel(HttpLogDetailLevel.BASIC))
-                        .credential(sharedKeyCred);
-
-        if (getTestMode() == TestMode.PLAYBACK) {
-            clientBuilderWithSharedKey
-                    .httpClient(interceptorManager.getPlaybackClient())
-                    .credential(request -> Mono.just(new AccessToken("this_is_a_token", OffsetDateTime.MAX)));
-        } else if (getTestMode() == TestMode.RECORD) {
-            clientBuilderWithSharedKey.addPolicy(interceptorManager.getRecordPolicy());
-        }
-
-        poolClientWithSharedKey = clientBuilderWithSharedKey.buildPoolClient();
+        batchClientBuilder.credential(sharedKeyCred);
+        poolClientWithSharedKey = batchClientBuilder.buildPoolClient();
     }
 
     @Test

diff --git a/...re/compute/batch/generated/TaskTests.java → ...va/com/azure/compute/batch/TaskTests.java b/...re/compute/batch/generated/TaskTests.java → ...va/com/azure/compute/batch/TaskTests.java
@@ -1,6 +1,5 @@
-package com.azure.compute.batch.generated;
+package com.azure.compute.batch;
 
-import com.azure.compute.batch.FileClient;
 import com.azure.compute.batch.models.*;
 import com.azure.core.http.rest.PagedIterable;
 import com.azure.core.util.BinaryData;
@@ -66,7 +65,12 @@ public void testJobUser() throws Exception {
             Assertions.assertNotNull(task);
             Assertions.assertEquals(taskId, task.getId());
             Assertions.assertEquals("test-user", task.getUserIdentity().getUsername());
-            Assertions.assertEquals("msmpi", task.getApplicationPackageReferences().get(0).getApplicationId());
+
+            //Recording file automatically sanitizes Application Id - Only verify App Id in Record Mode
+            if (getTestMode() == TestMode.RECORD) {
+                Assertions.assertEquals("msmpi", task.getApplicationPackageReferences().get(0).getApplicationId());
+            }
+
 
         } finally {
             try {
@@ -103,11 +107,14 @@ public void canCRUDTest() throws Exception {
 
             //The Storage operations run only in Record mode.
             // Playback mode is configured to test Batch operations only.
-            if (getTestMode() == TestMode.RECORD) {
+            if (getTestMode() != TestMode.PLAYBACK) {
                 // Create storage container
                 container = createBlobContainer(storageAccountName, storageAccountKey, "testingtaskcreate");
                 sas = uploadFileToCloud(container, BLOB_FILE_NAME, temp.getAbsolutePath());
             }
+            else {
+                sas = REDACTED;
+            }
 
             // Associate resource file with task
             ResourceFile file = new ResourceFile();
@@ -165,9 +172,12 @@ public void canCRUDTest() throws Exception {
 
                 //The Storage operations run only in Record mode.
                 // Playback mode is configured to test Batch operations only.
-                if (getTestMode() == TestMode.RECORD) {
+                if (getTestMode() != TestMode.PLAYBACK) {
                     outputSas = generateContainerSasToken(container);
                 }
+                else {
+                    outputSas = REDACTED;
+                }
                 // UPLOAD LOG
                 UploadBatchServiceLogsConfiguration logsConfiguration = new UploadBatchServiceLogsConfiguration(outputSas, OffsetDateTime.now().minusMinutes(-10));
                 UploadBatchServiceLogsResult uploadBatchServiceLogsResult = batchClientBuilder.buildComputeNodesClient().uploadBatchServiceLogs(liveIaasPoolId, task.getNodeInfo().getNodeId(), logsConfiguration);