mgmt, add more methods with context

sdk/resourcemanager/azure-resourcemanager-compute/CHANGELOG.md

@@ -10,6 +10,7 @@
 - Supported `beginDeleteByResourceGroup` and `beginDeleteById` in `Snapshots`.
 - Supported `beginCreate(Context)` in `Disk`.
 - Supported `beginDeleteByResourceGroup(Context)` and `beginDeleteById(Context)` in `Disks`.
+- Supported `getPrimaryNetworkInterface(Context)` in `VirtualMachine`.
 
 ### Breaking Changes
 

sdk/resourcemanager/azure-resourcemanager-compute/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "java",
   "TagPrefix": "java/resourcemanager/azure-resourcemanager-compute",
-  "Tag": "java/resourcemanager/azure-resourcemanager-compute_1b89db415f"
+  "Tag": "java/resourcemanager/azure-resourcemanager-compute_40571cd399"
 }

sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/implementation/VirtualMachineImpl.java

@@ -1735,6 +1735,11 @@ public NetworkInterface getPrimaryNetworkInterface() {
         return this.getPrimaryNetworkInterfaceAsync().block();
     }
 
+    @Override
+    public NetworkInterface getPrimaryNetworkInterface(Context context) {
+        return this.getPrimaryNetworkInterfaceAsync().contextWrite(c -> FluxUtil.toReactorContext(context)).block();
+    }
+
     @Override
     public Mono<NetworkInterface> getPrimaryNetworkInterfaceAsync() {
         return this.networkManager.networkInterfaces().getByIdAsync(primaryNetworkInterfaceId());

sdk/resourcemanager/azure-resourcemanager-compute/src/test/java/com/azure/resourcemanager/compute/VirtualMachineOperationsTests.java

@@ -6,6 +6,10 @@
 import com.azure.core.http.HttpMethod;
 import com.azure.core.http.HttpPipeline;
 import com.azure.core.http.HttpPipelineBuilder;
+import com.azure.core.http.HttpPipelineCallContext;
+import com.azure.core.http.HttpPipelineNextPolicy;
+import com.azure.core.http.HttpPipelinePosition;
+import com.azure.core.http.HttpResponse;
 import com.azure.core.http.policy.HttpPipelinePolicy;
 import com.azure.core.http.rest.PagedIterable;
 import com.azure.core.http.rest.PagedResponse;
@@ -83,6 +87,7 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Mono;
 
 import java.time.Duration;
 import java.util.ArrayList;
@@ -2311,6 +2316,7 @@ public void canCreateDiskWithShares() {
 
     @Test
     public void canBeginCreateAndDeleteWithContext() {
+        rgName = null;
         final String vmName = generateRandomResourceName("jvm", 15);
         final String diskName = generateRandomResourceName("jvdsk", 15);
         final String snapshotName = generateRandomResourceName("jvss", 15);
@@ -2319,45 +2325,61 @@ public void canBeginCreateAndDeleteWithContext() {
         final String publicIpDnsLabel = generateRandomResourceName("pip", 20);
         final AtomicInteger createCounter = new AtomicInteger(0);
         final AtomicInteger deleteCounter = new AtomicInteger(0);
-        HttpPipelinePolicy verificationPolicy = (context, next) -> {
-            if (context.getHttpRequest().getHttpMethod() == HttpMethod.PUT) {
-                // verify that all co-related resource creation requests will have the Context information
-                Object correlationData = context.getContext().getData(correlationKey).get();
-                Assertions.assertEquals(correlationId, correlationData);
-                createCounter.incrementAndGet();
-            } else if (context.getHttpRequest().getHttpMethod() == HttpMethod.DELETE) {
-                // verify that all co-related resource deletion requests will have the Context information
-                Object correlationData = context.getContext().getData(correlationKey).get();
-                Assertions.assertEquals(correlationId, correlationData);
-                deleteCounter.incrementAndGet();
+        final AtomicInteger getCounter = new AtomicInteger();
+        HttpPipelinePolicy verificationPolicy = new HttpPipelinePolicy() {
+            @Override
+            public HttpPipelinePosition getPipelinePosition() {
+                return HttpPipelinePosition.PER_CALL;
+            }
+
+            @Override
+            public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+
+                if (context.getHttpRequest().getHttpMethod() == HttpMethod.PUT) {
+                    // verify that all co-related resource creation requests will have the Context information
+                    Object correlationData = context.getContext().getData(correlationKey).get();
+                    Assertions.assertEquals(correlationId, correlationData);
+                    createCounter.incrementAndGet();
+                } else if (context.getHttpRequest().getHttpMethod() == HttpMethod.DELETE) {
+                    // verify that all co-related resource deletion requests will have the Context information
+                    Object correlationData = context.getContext().getData(correlationKey).get();
+                    Assertions.assertEquals(correlationId, correlationData);
+                    deleteCounter.incrementAndGet();
+                } else if (context.getHttpRequest().getHttpMethod() == HttpMethod.GET) {
+                    // some GET requests are nested inside implementations, thus only verify methods we are interested in
+                    context.getData(correlationKey).ifPresent(data -> {
+                        Assertions.assertEquals(correlationId, data);
+                        getCounter.incrementAndGet();
+                    });
+                }
+                return next.process();
             }
-            return next.process();
         };
-        ComputeManager localComputeManager
+        ComputeManager computeManagerWithPolicy
             = buildManager(ComputeManager.class, computeManager.httpPipeline(), verificationPolicy);
         Context context = new Context(correlationKey, correlationId);
 
-        Disk disk = localComputeManager.disks()
+        Disk disk = computeManagerWithPolicy.disks()
             .define(diskName)
             .withRegion(region)
-            .withNewResourceGroup(rgName)
+            .withNewResourceGroup(rgName2)
             .withData()
             .withSizeInGB(1)
             .beginCreate(context)
             .getFinalResult();
 
-        Snapshot snapshot = localComputeManager.snapshots()
+        Snapshot snapshot = computeManagerWithPolicy.snapshots()
             .define(snapshotName)
             .withRegion(region)
-            .withExistingResourceGroup(rgName)
+            .withExistingResourceGroup(rgName2)
             .withDataFromDisk(disk)
             .beginCreate(context)
             .getFinalResult();
 
-        Accepted<VirtualMachine> accepted = localComputeManager.virtualMachines()
+        Accepted<VirtualMachine> accepted = computeManagerWithPolicy.virtualMachines()
             .define(vmName)
             .withRegion(region)
-            .withExistingResourceGroup(rgName)
+            .withExistingResourceGroup(rgName2)
             .withNewPrimaryNetwork("10.0.0.0/28")
             .withPrimaryPrivateIPAddressDynamic()
             .withNewPrimaryPublicIPAddress(publicIpDnsLabel)
@@ -2370,18 +2392,22 @@ public void canBeginCreateAndDeleteWithContext() {
             .withPrimaryNetworkInterfaceDeleteOptions(DeleteOptions.DETACH)
             .beginCreate(context);
         VirtualMachine vm = accepted.getFinalResult();
-        String nicId = vm.getPrimaryNetworkInterface().id();
+
+        int getCountWithContext = getCounter.get();
+        String nicId = vm.getPrimaryNetworkInterface(context).id();
+        Assertions.assertEquals(1, getCounter.get() - getCountWithContext);
 
         // resourceGroup + disk + snapshot + network + neworkInterface + publicIp + VM = 7
         Assertions.assertEquals(7, createCounter.get());
 
-        localComputeManager.virtualMachines().beginDeleteById(vm.id(), false, context).getFinalResult();
-        localComputeManager.networkManager().networkInterfaces().beginDeleteById(nicId, context);
-        localComputeManager.snapshots().beginDeleteById(snapshot.id(), context).getFinalResult();
-        localComputeManager.disks().beginDeleteById(disk.id(), context).getFinalResult();
+        computeManagerWithPolicy.virtualMachines().beginDeleteById(vm.id(), false, context).getFinalResult();
+        computeManagerWithPolicy.networkManager().networkInterfaces().beginDeleteById(nicId, context);
+        computeManagerWithPolicy.snapshots().beginDeleteById(snapshot.id(), context).getFinalResult();
+        computeManagerWithPolicy.disks().beginDeleteById(disk.id(), context).getFinalResult();
+        computeManagerWithPolicy.resourceManager().resourceGroups().beginDeleteByName(rgName2, context);
 
-        // vm + nic + snapshot + disk = 4
-        Assertions.assertEquals(4, deleteCounter.get());
+        // vm + nic + snapshot + disk + resourcegroup = 5
+        Assertions.assertEquals(5, deleteCounter.get());
     }
 
     // *********************************** helper methods ***********************************
@@ -2391,7 +2417,7 @@ private <T> T buildManager(Class<T> managerClazz, HttpPipeline httpPipeline, Htt
             pipelinePolicies.add(httpPipeline.getPolicy(i));
         }
 
-        pipelinePolicies.add(policy);
+        pipelinePolicies.add(0, policy);
 
         HttpPipeline newPipeline = new HttpPipelineBuilder().httpClient(httpPipeline.getHttpClient())
             .policies(pipelinePolicies.toArray(new HttpPipelinePolicy[0]))

sdk/resourcemanager/azure-resourcemanager-network/CHANGELOG.md

@@ -4,7 +4,11 @@
 
 ### Features Added
 
-- Supported `beginDeleteByResourceGroup` and `beginDeleteById` in `NetworkInterfaces`.
+- Supported `listByResourceGroup(String, Context)`, `getPrimaryNetworkInterface(Context)`, `beginDeleteByResourceGroup` and `beginDeleteById` in `NetworkInterfaces`.
+- Supported `getByResourceGroup(String, String, Context)` in `ApplicationSecurityGroups`.
+- Supported `getByResourceGroup(String, String, Context)` in `NetworkSecurityGroups`.
+- Supported `getByResourceGroup(String, String, Context)` in `Networks`.
+- Supported `listAssociatedApplicationSecurityGroups(Context)` in `NicIpConfigurationBase`.
 
 ### Breaking Changes
 

sdk/resourcemanager/azure-resourcemanager-network/src/main/java/com/azure/resourcemanager/network/implementation/ApplicationSecurityGroupsImpl.java

@@ -2,6 +2,8 @@
 // Licensed under the MIT License.
 package com.azure.resourcemanager.network.implementation;
 
+import com.azure.core.util.Context;
+import com.azure.core.util.FluxUtil;
 import com.azure.resourcemanager.network.NetworkManager;
 import com.azure.resourcemanager.network.fluent.ApplicationSecurityGroupsClient;
 import com.azure.resourcemanager.network.fluent.models.ApplicationSecurityGroupInner;
@@ -36,4 +38,11 @@ protected ApplicationSecurityGroupImpl wrapModel(ApplicationSecurityGroupInner i
         }
         return new ApplicationSecurityGroupImpl(inner.name(), inner, this.manager());
     }
+
+    @Override
+    public ApplicationSecurityGroup getByResourceGroup(String resourceGroupName, String name, Context context) {
+        return getByResourceGroupAsync(resourceGroupName, name)
+            .contextWrite(c -> c.putAll(FluxUtil.toReactorContext(context).readOnly()))
+            .block();
+    }
 }

sdk/resourcemanager/azure-resourcemanager-network/src/main/java/com/azure/resourcemanager/network/implementation/NetworkInterfacesImpl.java

@@ -129,4 +129,12 @@ public Accepted<Void> beginDeleteByResourceGroup(String resourceGroupName, Strin
                 .block(),
             Function.identity(), Void.class, null, context);
     }
+
+    @Override
+    public PagedIterable<NetworkInterface> listByResourceGroup(String resourceGroupName, Context context) {
+        return this.inner()
+            .listByResourceGroup(resourceGroupName, context)
+            .mapPage(networkInterfaceInner -> new NetworkInterfaceImpl(networkInterfaceInner.name(),
+                networkInterfaceInner, NetworkInterfacesImpl.this.manager()));
+    }
 }

sdk/resourcemanager/azure-resourcemanager-network/src/main/java/com/azure/resourcemanager/network/implementation/NetworkSecurityGroupsImpl.java

@@ -2,7 +2,9 @@
 // Licensed under the MIT License.
 package com.azure.resourcemanager.network.implementation;
 
+import com.azure.core.util.Context;
 import com.azure.core.util.CoreUtils;
+import com.azure.core.util.FluxUtil;
 import com.azure.resourcemanager.network.NetworkManager;
 import com.azure.resourcemanager.network.fluent.NetworkSecurityGroupsClient;
 import com.azure.resourcemanager.network.fluent.models.NetworkSecurityGroupInner;
@@ -68,4 +70,11 @@ protected NetworkSecurityGroupImpl wrapModel(NetworkSecurityGroupInner inner) {
         }
         return new NetworkSecurityGroupImpl(inner.name(), inner, this.manager());
     }
+
+    @Override
+    public NetworkSecurityGroup getByResourceGroup(String resourceGroupName, String name, Context context) {
+        return getByResourceGroupAsync(resourceGroupName, name)
+            .contextWrite(c -> c.putAll(FluxUtil.toReactorContext(context).readOnly()))
+            .block();
+    }
 }

sdk/resourcemanager/azure-resourcemanager-network/src/main/java/com/azure/resourcemanager/network/implementation/NetworksImpl.java

@@ -2,6 +2,8 @@
 // Licensed under the MIT License.
 package com.azure.resourcemanager.network.implementation;
 
+import com.azure.core.util.Context;
+import com.azure.core.util.FluxUtil;
 import com.azure.resourcemanager.network.NetworkManager;
 import com.azure.resourcemanager.network.fluent.VirtualNetworksClient;
 import com.azure.resourcemanager.network.fluent.models.SubnetInner;
@@ -71,4 +73,11 @@ protected NetworkImpl wrapModel(VirtualNetworkInner inner) {
         }
         return new NetworkImpl(inner.name(), inner, this.manager());
     }
+
+    @Override
+    public Network getByResourceGroup(String resourceGroupName, String name, Context context) {
+        return getByResourceGroupAsync(resourceGroupName, name)
+            .contextWrite(c -> c.putAll(FluxUtil.toReactorContext(context).readOnly()))
+            .block();
+    }
 }

sdk/resourcemanager/azure-resourcemanager-network/src/main/java/com/azure/resourcemanager/network/implementation/NicIpConfigurationBaseImpl.java

@@ -4,7 +4,9 @@
 package com.azure.resourcemanager.network.implementation;
 
 import com.azure.core.management.SubResource;
+import com.azure.core.util.Context;
 import com.azure.core.util.CoreUtils;
+import com.azure.core.util.FluxUtil;
 import com.azure.resourcemanager.network.NetworkManager;
 import com.azure.resourcemanager.network.fluent.models.ApplicationSecurityGroupInner;
 import com.azure.resourcemanager.network.models.ApplicationGatewayBackend;
@@ -171,13 +173,20 @@ public List<LoadBalancerInboundNatRule> listAssociatedLoadBalancerInboundNatRule
 
     @Override
     public List<ApplicationSecurityGroup> listAssociatedApplicationSecurityGroups() {
+        return listAssociatedApplicationSecurityGroups(Context.NONE);
+    }
+
+    @Override
+    public List<ApplicationSecurityGroup> listAssociatedApplicationSecurityGroups(Context context) {
         if (CoreUtils.isNullOrEmpty(this.innerModel().applicationSecurityGroups())) {
             return Collections.emptyList();
         }
 
         List<ApplicationSecurityGroup> applicationSecurityGroups = Flux
             .fromStream(this.innerModel().applicationSecurityGroups().stream().map(ApplicationSecurityGroupInner::id))
-            .flatMapSequential(id -> this.networkManager.applicationSecurityGroups().getByIdAsync(id))
+            .flatMapSequential(id -> this.networkManager.applicationSecurityGroups()
+                .getByIdAsync(id)
+                .contextWrite(c -> c.putAll(FluxUtil.toReactorContext(context).readOnly())))
             .collectList()
             .block();
         return applicationSecurityGroups == null

sdk/resourcemanager/azure-resourcemanager-network/src/main/java/com/azure/resourcemanager/network/models/ApplicationSecurityGroups.java

@@ -3,6 +3,7 @@
 package com.azure.resourcemanager.network.models;
 
 import com.azure.core.annotation.Fluent;
+import com.azure.core.util.Context;
 import com.azure.resourcemanager.network.NetworkManager;
 import com.azure.resourcemanager.resources.fluentcore.arm.collection.SupportsBatchDeletion;
 import com.azure.resourcemanager.resources.fluentcore.arm.collection.SupportsDeletingByResourceGroup;
@@ -22,4 +23,14 @@ public interface ApplicationSecurityGroups extends SupportsCreating<ApplicationS
     SupportsGettingByResourceGroup<ApplicationSecurityGroup>, SupportsGettingById<ApplicationSecurityGroup>,
     SupportsDeletingById, SupportsDeletingByResourceGroup, SupportsBatchCreation<ApplicationSecurityGroup>,
     SupportsBatchDeletion, HasManager<NetworkManager> {
+
+    /**
+     * Gets the information about a resource from Azure based on the resource name and the name of its resource group.
+     *
+     * @param resourceGroupName the name of the resource group the resource is in
+     * @param name the name of the application security group. (Note, this is not the ID)
+     * @param context the {@link Context} of the request
+     * @return an immutable representation of the resource
+     */
+    ApplicationSecurityGroup getByResourceGroup(String resourceGroupName, String name, Context context);
 }

sdk/resourcemanager/azure-resourcemanager-network/src/main/java/com/azure/resourcemanager/network/models/HasNetworkInterfaces.java

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 package com.azure.resourcemanager.network.models;
 
+import com.azure.core.util.Context;
 import com.azure.resourcemanager.resources.fluentcore.arm.models.HasId;
 import reactor.core.publisher.Mono;
 
@@ -18,6 +19,16 @@ public interface HasNetworkInterfaces extends HasId {
      */
     NetworkInterface getPrimaryNetworkInterface();
 
+    /**
+     * Gets the primary network interface.
+     *
+     * <p>Note that this method can result in a call to the cloud to fetch the network interface information.
+     *
+     * @param context the {@link Context} of the request
+     * @return the primary network interface associated with this resource
+     */
+    NetworkInterface getPrimaryNetworkInterface(Context context);
+
     /**
      * Gets the primary network interface.
      *

sdk/resourcemanager/azure-resourcemanager-network/src/main/java/com/azure/resourcemanager/network/models/NetworkInterfaces.java

@@ -128,4 +128,13 @@ PagedIterable<VirtualMachineScaleSetNetworkInterface> listByVirtualMachineScaleS
      * @return the accepted deleting operation
      */
     Accepted<Void> beginDeleteByResourceGroup(String resourceGroupName, String name, Context context);
+
+    /**
+     * Lists resources of the specified type in the specified resource group.
+     *
+     * @param resourceGroupName the name of the resource group to list the resources from
+     * @param context the {@link Context} of the request
+     * @return the list of resources
+     */
+    PagedIterable<NetworkInterface> listByResourceGroup(String resourceGroupName, Context context);
 }

sdk/resourcemanager/azure-resourcemanager-network/src/main/java/com/azure/resourcemanager/network/models/NetworkSecurityGroups.java

@@ -3,6 +3,7 @@
 package com.azure.resourcemanager.network.models;
 
 import com.azure.core.annotation.Fluent;
+import com.azure.core.util.Context;
 import com.azure.resourcemanager.network.NetworkManager;
 import com.azure.resourcemanager.resources.fluentcore.arm.collection.SupportsBatchDeletion;
 import com.azure.resourcemanager.resources.fluentcore.arm.collection.SupportsDeletingByResourceGroup;
@@ -22,4 +23,13 @@ public interface NetworkSecurityGroups
     SupportsListingByResourceGroup<NetworkSecurityGroup>, SupportsGettingByResourceGroup<NetworkSecurityGroup>,
     SupportsGettingById<NetworkSecurityGroup>, SupportsDeletingById, SupportsDeletingByResourceGroup,
     SupportsBatchCreation<NetworkSecurityGroup>, SupportsBatchDeletion, HasManager<NetworkManager> {
+    /**
+     * Gets the information about {@link NetworkSecurityGroup} based on the resource name and the name of its resource group.
+     *
+     * @param resourceGroupName the name of the resource group the resource is in
+     * @param name the name of the network security group. (Note, this is not the ID)
+     * @param context the {@link Context} of the request
+     * @return an immutable representation of the resource
+     */
+    NetworkSecurityGroup getByResourceGroup(String resourceGroupName, String name, Context context);
 }