Change parsing logic in ATC custom cmdlet (Azure#24828)

* Change parsing logic in ATC custom cmdlet * Add if check * add if check * Minor change * Add UT
qinzhouxu · Jun 5, 2024 · fe81647 · fe81647
1 parent b108f01
commit fe81647
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 29 deletions.
diff --git a/src/NetworkFunction/NetworkFunction.Autorest/custom/New-AzNetworkFunctionCollectorPolicy.ps1 b/src/NetworkFunction/NetworkFunction.Autorest/custom/New-AzNetworkFunctionCollectorPolicy.ps1
@@ -167,23 +167,33 @@ param(
 )
 
 process {
-      $rg = $PSBoundParameters.ResourceGroupName
-
+      $hasIngestionPolicyIngestionSource = $PSBoundParameters.Remove('IngestionPolicyIngestionSource')
       # Ensure exr circuit bandwidth 1G or more
-      $cktname = $IngestionPolicyIngestionSource.ResourceId | Where {$IngestionPolicyIngestionSource.ResourceId -match "/*subscriptions/(?<subid>.*)/resourceGroups/(?<rgname>.*)/providers/Microsoft.Network/expressRouteCircuits/(?<circuitname>.*)"} | Foreach {$Matches['circuitname']}
-      Import-Module Az.Network -Force
-      $exrCircuit = Get-AzExpressRouteCircuit -Name $cktname -ResourceGroupName $rg
-      $bandwidthInGbps = $exrCircuit.BandwidthInGbps
-      $bandwidthInMbps = $exrCircuit.ServiceProviderProperties.BandwidthInMbps
-
-      if ($bandwidthInGbps -and ($bandwidthInGbps -lt 1)) {
-        throw "CollectorPolicy can not be updated because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
-      }
-
-      if ($bandwidthInMbps -and ($bandwidthInMbps -lt 1000)) {
-        throw "CollectorPolicy can not be updated because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
+      if ($hasIngestionPolicyIngestionSource) {
+        Import-Module Az.Network -Force
+        $ResourceIdSplit = $IngestionPolicyIngestionSource.ResourceId.Split(' ')
+        foreach ($ResourceId in $ResourceIdSplit)
+        {
+          $Splits = $ResourceId -split "/"
+          $cktsub = $Splits[2]
+          $cktrg = $Splits[4]
+          $cktname = $Splits[8]
+          Set-AzContext $cktsub -ErrorVariable notPresent -ErrorAction SilentlyContinue
+          $exrCircuit = Get-AzExpressRouteCircuit -Name $cktname -ResourceGroupName $cktrg
+          $bandwidthInGbps = $exrCircuit.BandwidthInGbps
+          $bandwidthInMbps = $exrCircuit.ServiceProviderProperties.BandwidthInMbps
+
+          if ($bandwidthInGbps -and ($bandwidthInGbps -lt 1)) {
+            throw "CollectorPolicy can not be created because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
+          }
+
+          if ($bandwidthInMbps -and ($bandwidthInMbps -lt 1000)) {
+            throw "CollectorPolicy can not be created because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
+          }
+        }
       }
-
+      Set-AzContext $SubscriptionId -ErrorVariable notPresent -ErrorAction SilentlyContinue
+      $PSBoundParameters.Add('IngestionPolicyIngestionSource', $IngestionPolicyIngestionSource)
       Az.NetworkFunction.internal\New-AzNetworkFunctionCollectorPolicy @PSBoundParameters
     }
 }
diff --git a/...tworkFunction/NetworkFunction.Autorest/custom/Update-AzNetworkFunctionCollectorPolicy.ps1 b/...tworkFunction/NetworkFunction.Autorest/custom/Update-AzNetworkFunctionCollectorPolicy.ps1
@@ -173,23 +173,33 @@ process {
       $null = $PSBoundParameters.Remove('WhatIf')
       $null = $PSBoundParameters.Remove('Confirm')
       $null = $PSBoundParameters.Remove('Location')
-      $rg = $PSBoundParameters.ResourceGroupName
 
       # 2. Ensure exr circuit bandwidth 1G or more
-      $cktname = $IngestionPolicyIngestionSource.ResourceId | Where {$IngestionPolicyIngestionSource.ResourceId -match "/*subscriptions/(?<subid>.*)/resourceGroups/(?<rgname>.*)/providers/Microsoft.Network/expressRouteCircuits/(?<circuitname>.*)"} | Foreach {$Matches['circuitname']}
-      Import-Module Az.Network -Force
-      $exrCircuit = Get-AzExpressRouteCircuit -Name $cktname -ResourceGroupName $rg
-      $bandwidthInGbps = $exrCircuit.BandwidthInGbps
-      $bandwidthInMbps = $exrCircuit.ServiceProviderProperties.BandwidthInMbps
-
-      if ($bandwidthInGbps -and ($bandwidthInGbps -lt 1)) {
-        throw "CollectorPolicy can not be updated because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
-      }
-
-      if ($bandwidthInMbps -and ($bandwidthInMbps -lt 1000)) {
-        throw "CollectorPolicy can not be updated because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
+      if ($hasIngestionPolicyIngestionSource) {
+        Import-Module Az.Network -Force
+        # Ensure exr circuit bandwidth 1G or more
+        $ResourceIdSplit = $IngestionPolicyIngestionSource.ResourceId.Split(' ')
+        foreach ($ResourceId in $ResourceIdSplit)
+        {
+          $Splits = $ResourceId -split "/"
+          $cktsub = $Splits[2]
+          $cktrg = $Splits[4]
+          $cktname = $Splits[8]
+          Set-AzContext $cktsub -ErrorVariable notPresent -ErrorAction SilentlyContinue
+          $exrCircuit = Get-AzExpressRouteCircuit -Name $cktname -ResourceGroupName $cktrg
+          $bandwidthInGbps = $exrCircuit.BandwidthInGbps
+          $bandwidthInMbps = $exrCircuit.ServiceProviderProperties.BandwidthInMbps
+
+          if ($bandwidthInGbps -and ($bandwidthInGbps -lt 1)) {
+            throw "CollectorPolicy can not be updated because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
+          }
+
+          if ($bandwidthInMbps -and ($bandwidthInMbps -lt 1000)) {
+            throw "CollectorPolicy can not be updated because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
+          }
+        }
       }
-
+      Set-AzContext $SubscriptionId -ErrorVariable notPresent -ErrorAction SilentlyContinue
       $cp = Get-AzNetworkFunctionCollectorPolicy @PSBoundParameters
 
       # 3. PUT

diff --git a/...workFunction/NetworkFunction.Autorest/test/New-AzNetworkFunctionCollectorPolicy.Tests.ps1 b/...workFunction/NetworkFunction.Autorest/test/New-AzNetworkFunctionCollectorPolicy.Tests.ps1
@@ -27,6 +27,12 @@ Describe 'New-AzNetworkFunctionCollectorPolicy' {
         } 
     }
 
+    It 'CreateExpanded2' {
+        { 
+            { New-AzNetworkFunctionCollectorPolicy -collectorpolicyname $env.collectorPolicyName -azuretrafficcollectorname $env.azureTrafficCollectorName -resourcegroupname $env.resourceGroup -location $env.location -IngestionPolicyIngestionSource @{ResourceId = $env.ResourceIdLessThan1G}, @{ResourceId = $env.ResourceId1G} -IngestionPolicyIngestionType $env.IngestionType } | Should -Throw -ExpectedMessage "CollectorPolicy can not be updated because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
+        } 
+    }
+
     It 'Create' {
         { 
             { New-AzNetworkFunctionCollectorPolicy -collectorpolicyname $env.collectorPolicyName -azuretrafficcollectorname $env.azureTrafficCollectorName -resourcegroupname $env.resourceGroup -location $env.location -IngestionPolicyIngestionSource @{ResourceId = $env.ResourceId1G} -IngestionPolicyIngestionType $env.IngestionType } | Should -Not -Throw

diff --git a/...kFunction/NetworkFunction.Autorest/test/Update-AzNetworkFunctionCollectorPolicy.Tests.ps1 b/...kFunction/NetworkFunction.Autorest/test/Update-AzNetworkFunctionCollectorPolicy.Tests.ps1
@@ -26,4 +26,10 @@ Describe 'Update-AzNetworkFunctionCollectorPolicy' {
             { Update-AzNetworkFunctionCollectorPolicy -collectorpolicyname $env.collectorPolicyName -azuretrafficcollectorname $env.azureTrafficCollectorName -resourcegroupname $env.resourceGroup -location $env.location -IngestionPolicyIngestionSource @{ResourceId = $env.ResourceIdLessThan1G} -IngestionPolicyIngestionType $env.IngestionType } | Should Throw -ExpectedMessage "CollectorPolicy can not be updated because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
         }
     }
+
+    It 'UpdateExpanded2' {
+        { 
+            { Update-AzNetworkFunctionCollectorPolicy -collectorpolicyname $env.collectorPolicyName -azuretrafficcollectorname $env.azureTrafficCollectorName -resourcegroupname $env.resourceGroup -location $env.location -IngestionPolicyIngestionSource @{ResourceId = $env.ResourceIdLessThan1G}, @{ResourceId = $env.ResourceId1G} -IngestionPolicyIngestionType $env.IngestionType } | Should Throw -ExpectedMessage "CollectorPolicy can not be updated because circuit has bandwidth less than 1G. Circuit size with a bandwidth of 1G or more is supported."
+        }
+    }
 }