Support CosmosDb for PrivateEndpointConnection and PrivateLinkResource operations

src/CosmosDB/CosmosDB.Test/CosmosDB.Test.csproj

@@ -5,5 +5,6 @@
   <Import Project="$(MSBuildThisFileDirectory)..\..\Az.Test.props" />
   <ItemGroup>
     <PackageReference Include="Microsoft.Azure.Management.CosmosDB" Version="1.1.0" />
+    <PackageReference Include="Microsoft.Azure.Management.Network" Version="19.19.0-preview" />
 </ItemGroup>
 </Project>

src/CosmosDB/CosmosDB.Test/ScenarioTests/AccountTests.cs

@@ -54,5 +54,12 @@ public void TestAddRegionOperation()
         {
             TestController.NewInstance.RunPowerShellTest(_logger, "Test-AddRegionOperation");
         }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void TestPrivateEndpoint()
+        {
+            TestController.NewInstance.RunPowerShellTest(_logger, "Test-PrivateEndpoint");
+        }
     }
 }

src/CosmosDB/CosmosDB.Test/ScenarioTests/AccountTests.ps1

@@ -188,4 +188,49 @@ function Test-AddRegionOperation
   $updatedCosmosDBAccount = Update-AzCosmosDBAccountRegion -ResourceGroupName $rgName -Name $cosmosDBAccountName -Location $locationlist
   $updatedCosmosDBAccount = Get-AzCosmosDBAccount -ResourceGroupName $rgName -Name $cosmosDBAccountName
   Assert-AreEqual $cosmosDBAccount.Locations.Count $updatedCosmosDBAccount.Locations.Count - 1 
-}
+}
+
+function Test-PrivateEndpoint
+{
+  # Setup
+  $location = "East US"
+  $peName = "mype";
+  $storageAccount = "xdmsa2";
+
+  #use an existing account with the following properties
+  $cosmosDBExistingAccountName = "db9934121" 
+  $rgname = "CosmosDBResourceGroup9507"
+
+  $cosmosDBAccount = Get-AzCosmosDBAccount -ResourceGroupName $rgname -Name $cosmosDBExistingAccountName
+  $resourceId = $cosmosDBAccount.Id
+
+  $peSubnet = New-AzVirtualNetworkSubnetConfig -Name peSubnet -AddressPrefix "11.0.1.0/24" -PrivateEndpointNetworkPolicies "Disabled"
+  $vnetPE = New-AzVirtualNetwork -Name "vnetPE" -ResourceGroupName $rgname -Location $location -AddressPrefix "11.0.0.0/16" -Subnet $peSubnet
+
+  $plsConnection= New-AzPrivateLinkServiceConnection -Name plsConnection -PrivateLinkServiceId  $resourceId -GroupId 'Sql'
+  $privateEndpoint = New-AzPrivateEndpoint -ResourceGroupName $rgname -Name $peName -Location $location -Subnet $vnetPE.subnets[0] -PrivateLinkServiceConnection $plsConnection -ByManualRequest
+
+  $pecGet = Get-AzPrivateEndpointConnection -PrivateLinkResourceId $resourceId
+  Assert-NotNull $pecGet;
+  Assert-AreEqual "Pending" $pecGet.PrivateLinkServiceConnectionState.Status
+
+  # Approve Private Endpoint Connection
+  $pecApprove = Approve-AzPrivateEndpointConnection -ResourceId $pecGet.Id
+  Assert-NotNull $pecApprove;
+
+  $pecGet2 = Get-AzPrivateEndpointConnection -PrivateLinkResourceId $resourceId
+  do 
+  {
+    $pecGet2 = Get-AzPrivateEndpointConnection -PrivateLinkResourceId $resourceId
+  } while ($pecGet2.PrivateLinkServiceConnectionState.Status -ne "Approved")
+
+  # Remove Private Endpoint Connection
+  $pecRemove = Remove-AzPrivateEndpointConnection -ResourceId $pecGet.Id -PassThru -Force
+  Assert-AreEqual true $pecRemove
+
+  $pecGet3 = Get-AzPrivateEndpointConnection -PrivateLinkResourceId $resourceId
+  do 
+  {
+    $pecGet3 = Get-AzPrivateEndpointConnection -PrivateLinkResourceId $resourceId
+  } while (($pecGet3) -ne $null)
+}

src/CosmosDB/CosmosDB.Test/ScenarioTests/TestController.cs

@@ -18,8 +18,10 @@
 using System.IO;
 using System.Linq;
 using Microsoft.Azure.Commands.Common.Authentication;
+using Microsoft.Azure.Internal.Common;
 using Microsoft.Azure.Management.Internal.Resources;
 using Microsoft.Azure.Management.CosmosDB;
+using Microsoft.Azure.Management.Network;
 using Microsoft.Azure.Test.HttpRecorder;
 using Microsoft.Rest.ClientRuntime.Azure.TestFramework;
 using Microsoft.WindowsAzure.Commands.ScenarioTest;
@@ -36,6 +38,10 @@ public class TestController : RMTestBase
 
         public CosmosDBManagementClient CosmosDBManagementClient { get; private set; }
 
+        public NetworkManagementClient NetworkManagementClient { get; private set; }
+
+        public AzureRestClient AzureRestClient { get; private set; }
+
         public static TestController NewInstance => new TestController();
 
         protected TestController()
@@ -47,9 +53,13 @@ protected void SetupManagementClients(MockContext context)
         {
             ResourceManagementClient = GetResourceManagementClient(context);
             CosmosDBManagementClient = GetCosmosDBManagementClient(context);
+            NetworkManagementClient = GetNetworkManagementClient(context);
+            AzureRestClient = GetAzureRestClient(context);
             _helper.SetupManagementClients(
                 ResourceManagementClient,
-                CosmosDBManagementClient);
+                CosmosDBManagementClient,
+                NetworkManagementClient,
+                AzureRestClient);
         }
 
         public void RunPowerShellTest(ServiceManagement.Common.Models.XunitTracingInterceptor logger, params string[] scripts)
@@ -102,6 +112,7 @@ public void RunPsTestWorkflow(
                     "ScenarioTests\\" + callingClassName + ".ps1",
                     _helper.RMProfileModule,
                     _helper.GetRMModulePath("Az.CosmosDB.psd1"),
+                    _helper.GetRMModulePath("Az.Network.psd1"),
                     "AzureRM.Resources.ps1");
                 try
                 {
@@ -127,5 +138,15 @@ private static CosmosDBManagementClient GetCosmosDBManagementClient(MockContext
         {
             return context.GetServiceClient<CosmosDBManagementClient>(TestEnvironmentFactory.GetTestEnvironment());
         }
+
+        private static NetworkManagementClient GetNetworkManagementClient(MockContext context)
+        {
+            return context.GetServiceClient<NetworkManagementClient>(TestEnvironmentFactory.GetTestEnvironment());
+        }
+
+        private static AzureRestClient GetAzureRestClient(MockContext context)
+        {
+            return context.GetServiceClient<AzureRestClient>(TestEnvironmentFactory.GetTestEnvironment());
+        }
     }
 }