Add function in psm1 to mock subscription id (#1318)

* do not remove loadEnv.ps1 while regenerate/rebuild * add function get-localsubscription * rename function * adjust subscription id in export cmdlets * typo * remove environment variable * move get-subscriptionidtestsafe to utils * add template loadenv back * fix mockpipeline type * fix * add tenant back in parameter default
Azure · Mar 20, 2024 · 7359d8e · 7359d8e
1 parent c33b901
commit 7359d8e
Show file tree

Hide file tree

Showing 6 changed files with 141 additions and 89 deletions.
diff --git a/powershell/generators/psm1.ts b/powershell/generators/psm1.ts
@@ -149,5 +149,4 @@ ${getProfileExportScript(`Join-Path $PSScriptRoot '${project.exportsFolder}'`, p
   Write-Information "Loaded Module '$($instance.Name)'"`);
   psm1.trim();
   project.state.writeFile(project.psm1, `${psm1}`, undefined, 'source-file-powershell');
-}
-
+}
diff --git a/powershell/resources/assets/test-module.ps1 b/powershell/resources/assets/test-module.ps1
@@ -63,13 +63,14 @@ try
   if ($TestMode -ne 'playback')
   {
     setupEnv
+  } else {
+    $env:AzPSAutorestTestPlaybackMode = $true
   }
   $testFolder = Join-Path $PSScriptRoot '${$lib.path.relative($project.baseFolder, $project.testFolder)}'
   if ($null -ne $TestName)
   {
     Invoke-Pester -Script @{ Path = $testFolder } -TestName $TestName -ExcludeTag $ExcludeTag -EnableExit -OutputFile (Join-Path $testFolder "$moduleName-TestResults.xml")
-  } else
-  {
+  } else {
     Invoke-Pester -Script @{ Path = $testFolder } -ExcludeTag $ExcludeTag -EnableExit -OutputFile (Join-Path $testFolder "$moduleName-TestResults.xml")
   }
 } Finally
@@ -78,6 +79,9 @@ try
   {
     cleanupEnv
   }
+  else {
+    $env:AzPSAutorestTestPlaybackMode = ''
+  }
 }
 
 Write-Host -ForegroundColor Green '-------------Done-------------'
diff --git a/powershell/resources/assets/test/loadEnv.ps1 b/powershell/resources/assets/test/loadEnv.ps1
@@ -14,5 +14,5 @@ if (Test-Path -Path (Join-Path $PSScriptRoot $envFile)) {
 $env = @{}
 if (Test-Path -Path $envFilePath) {
     $env = Get-Content (Join-Path $PSScriptRoot $envFile) | ConvertFrom-Json
-    $PSDefaultParameterValues=@{"*:SubscriptionId"=$env.SubscriptionId; "*:Tenant"=$env.Tenant}
+    $PSDefaultParameterValues=@{"*:Tenant"=$env.Tenant}
 }
diff --git a/powershell/resources/psruntime/BuildTime/Cmdlets/ExportTestStub.cs b/powershell/resources/psruntime/BuildTime/Cmdlets/ExportTestStub.cs
@@ -12,44 +12,65 @@
 
 namespace Microsoft.Rest.ClientRuntime.PowerShell
 {
-  [Cmdlet(VerbsData.Export, "TestStub")]
-  [DoNotExport]
-  public class ExportTestStub : PSCmdlet
-  {
-    [Parameter(Mandatory = true)]
-    [ValidateNotNullOrEmpty]
-    public string ModuleName { get; set; }
-
-    [Parameter(Mandatory = true)]
-    [ValidateNotNullOrEmpty]
-    public string ExportsFolder { get; set; }
-
-    [Parameter(Mandatory = true)]
-    [ValidateNotNullOrEmpty]
-    public string OutputFolder { get; set; }
-
-    [Parameter]
-    public SwitchParameter IncludeGenerated { get; set; }
-
-    protected override void ProcessRecord()
+    [Cmdlet(VerbsData.Export, "TestStub")]
+    [DoNotExport]
+    public class ExportTestStub : PSCmdlet
     {
-      try
-      {
-        if (!Directory.Exists(ExportsFolder))
-        {
-          throw new ArgumentException($"Exports folder '{ExportsFolder}' does not exist");
-        }
+        [Parameter(Mandatory = true)]
+        [ValidateNotNullOrEmpty]
+        public string ModuleName { get; set; }
 
-        var exportDirectories = Directory.GetDirectories(ExportsFolder);
-        if (!exportDirectories.Any())
-        {
-          exportDirectories = new[] { ExportsFolder };
-        }
-        var utilFile = Path.Combine(OutputFolder, "utils.ps1");
-        if (!File.Exists(utilFile))
+        [Parameter(Mandatory = true)]
+        [ValidateNotNullOrEmpty]
+        public string ExportsFolder { get; set; }
+
+        [Parameter(Mandatory = true)]
+        [ValidateNotNullOrEmpty]
+        public string OutputFolder { get; set; }
+
+        [Parameter]
+        public SwitchParameter IncludeGenerated { get; set; }
+
+        protected override void ProcessRecord()
         {
-          var sc = new StringBuilder();
-          sc.AppendLine(@"function RandomString([bool]$allChars, [int32]$len) {
+            try
+            {
+                if (!Directory.Exists(ExportsFolder))
+                {
+                    throw new ArgumentException($"Exports folder '{ExportsFolder}' does not exist");
+                }
+
+                var exportDirectories = Directory.GetDirectories(ExportsFolder);
+                if (!exportDirectories.Any())
+                {
+                    exportDirectories = new[] { ExportsFolder };
+                }
+                /*var loadEnvFile = Path.Combine(OutputFolder, "loadEnv.ps1");
+                if (!File.Exists(loadEnvFile))
+                {
+                    var sc = new StringBuilder();
+                    sc.AppendLine(@"
+$envFile = 'env.json'
+if ($TestMode -eq 'live') {
+    $envFile = 'localEnv.json'
+}
+
+if (Test-Path -Path (Join-Path $PSScriptRoot $envFile)) {
+    $envFilePath = Join-Path $PSScriptRoot $envFile
+} else {
+    $envFilePath = Join-Path $PSScriptRoot '..\$envFile'
+}
+$env = @{}
+if (Test-Path -Path $envFilePath) {
+    $env = Get-Content (Join-Path $PSScriptRoot $envFile) | ConvertFrom-Json
+}");
+                    File.WriteAllText(loadEnvFile, sc.ToString());
+                }*/
+                var utilFile = Path.Combine(OutputFolder, "utils.ps1");
+                if (!File.Exists(utilFile))
+                {
+                    var sc = new StringBuilder();
+                    sc.AppendLine(@"function RandomString([bool]$allChars, [int32]$len) {
     if ($allChars) {
         return -join ((33..126) | Get-Random -Count $len | % {[char]$_})
     } else {
@@ -105,37 +126,40 @@ protected override void ProcessRecord()
     # Clean resources you create for testing
 }
 ");
-          File.WriteAllText(utilFile, sc.ToString());
-        }
-        foreach (var exportDirectory in exportDirectories)
-        {
-          var outputFolder = OutputFolder;
-          if (exportDirectory != ExportsFolder)
-          {
-            outputFolder = Path.Combine(OutputFolder, Path.GetFileName(exportDirectory));
-            Directory.CreateDirectory(outputFolder);
-          }
-
-          var variantGroups = GetScriptCmdlets(exportDirectory)
-              .SelectMany(fi => fi.ToVariants())
-              .Where(v => !v.IsDoNotExport)
-              .GroupBy(v => v.CmdletName)
-              .Select(vg => new VariantGroup(ModuleName, vg.Key, vg.Select(v => v).ToArray(), outputFolder, isTest: true))
-              .Where(vtg => !File.Exists(vtg.FilePath) && (IncludeGenerated || !vtg.IsGenerated));
-
-          foreach (var variantGroup in variantGroups)
-          {
-            var sb = new StringBuilder();
-            sb.AppendLine($"if(($null -eq $TestName) -or ($TestName -contains '{variantGroup.CmdletName}'))");
-            sb.AppendLine(@"{
+                    File.WriteAllText(utilFile, sc.ToString());
+                }
+
+
+
+                foreach (var exportDirectory in exportDirectories)
+                {
+                    var outputFolder = OutputFolder;
+                    if (exportDirectory != ExportsFolder)
+                    {
+                        outputFolder = Path.Combine(OutputFolder, Path.GetFileName(exportDirectory));
+                        Directory.CreateDirectory(outputFolder);
+                    }
+
+                    var variantGroups = GetScriptCmdlets(exportDirectory)
+                        .SelectMany(fi => fi.ToVariants())
+                        .Where(v => !v.IsDoNotExport)
+                        .GroupBy(v => v.CmdletName)
+                        .Select(vg => new VariantGroup(ModuleName, vg.Key, vg.Select(v => v).ToArray(), outputFolder, isTest: true))
+                        .Where(vtg => !File.Exists(vtg.FilePath) && (IncludeGenerated || !vtg.IsGenerated));
+
+                    foreach (var variantGroup in variantGroups)
+                    {
+                        var sb = new StringBuilder();
+                        sb.AppendLine($"if(($null -eq $TestName) -or ($TestName -contains '{variantGroup.CmdletName}'))");
+                        sb.AppendLine(@"{
   $loadEnvPath = Join-Path $PSScriptRoot 'loadEnv.ps1'
   if (-Not (Test-Path -Path $loadEnvPath)) {
       $loadEnvPath = Join-Path $PSScriptRoot '..\loadEnv.ps1'
   }
   . ($loadEnvPath)"
-);
-            sb.AppendLine($@"  $TestRecordingFile = Join-Path $PSScriptRoot '{variantGroup.CmdletName}.Recording.json'");
-            sb.AppendLine(@"  $currentPath = $PSScriptRoot
+            );
+                        sb.AppendLine($@"  $TestRecordingFile = Join-Path $PSScriptRoot '{variantGroup.CmdletName}.Recording.json'");
+                        sb.AppendLine(@"  $currentPath = $PSScriptRoot
   while(-not $mockingPath) {
       $mockingPath = Get-ChildItem -Path $currentPath -Recurse -Include 'HttpPipelineMocking.ps1' -File
       $currentPath = Split-Path -Path $currentPath -Parent
@@ -145,29 +169,29 @@ protected override void ProcessRecord()
 ");
 
 
-            sb.AppendLine($"Describe '{variantGroup.CmdletName}' {{");
-            var variants = variantGroup.Variants
-                .Where(v => IncludeGenerated || !v.Attributes.OfType<GeneratedAttribute>().Any())
-                .ToList();
+                        sb.AppendLine($"Describe '{variantGroup.CmdletName}' {{");
+                        var variants = variantGroup.Variants
+                            .Where(v => IncludeGenerated || !v.Attributes.OfType<GeneratedAttribute>().Any())
+                            .ToList();
+
+                        foreach (var variant in variants)
+                        {
+                            sb.AppendLine($"{Indent}It '{variant.VariantName}' -skip {{");
+                            sb.AppendLine($"{Indent}{Indent}{{ throw [System.NotImplementedException] }} | Should -Not -Throw");
+                            var variantSeparator = variants.IndexOf(variant) == variants.Count - 1 ? String.Empty : Environment.NewLine;
+                            sb.AppendLine($"{Indent}}}{variantSeparator}");
+                        }
+                        sb.AppendLine("}");
 
-            foreach (var variant in variants)
+                        File.WriteAllText(variantGroup.FilePath, sb.ToString());
+                    }
+                }
+            }
+            catch (Exception ee)
             {
-              sb.AppendLine($"{Indent}It '{variant.VariantName}' -skip {{");
-              sb.AppendLine($"{Indent}{Indent}{{ throw [System.NotImplementedException] }} | Should -Not -Throw");
-              var variantSeparator = variants.IndexOf(variant) == variants.Count - 1 ? String.Empty : Environment.NewLine;
-              sb.AppendLine($"{Indent}}}{variantSeparator}");
+                Console.WriteLine($"${ee.GetType().Name}/{ee.StackTrace}");
+                throw ee;
             }
-            sb.AppendLine("}");
-
-            File.WriteAllText(variantGroup.FilePath, sb.ToString());
-          }
         }
-      }
-      catch (Exception ee)
-      {
-        Console.WriteLine($"${ee.GetType().Name}/{ee.StackTrace}");
-        throw ee;
-      }
     }
-  }
 }
diff --git a/powershell/resources/psruntime/BuildTime/Models/PsProxyOutputs.cs b/powershell/resources/psruntime/BuildTime/Models/PsProxyOutputs.cs
@@ -290,9 +290,27 @@ private string GetDefaultValuesStatements()
                 {
                     setCondition = $" -and {defaultInfo.SetCondition}";
                 }
-                sb.AppendLine($"{Indent}{Indent}if (({variantListString}) -contains $parameterSet -and -not $PSBoundParameters.ContainsKey('{parameterName}'){setCondition}) {{");
-                sb.AppendLine($"{Indent}{Indent}{Indent}$PSBoundParameters['{parameterName}'] = {defaultInfo.Script}");
-                sb.Append($"{Indent}{Indent}}}");
+                //Yabo: this is bad to hard code the subscription id, but autorest load input README.md reversely (entry readme -> required readme), there are no other way to 
+                //override default value set in required readme
+                if ("SubscriptionId".Equals(parameterName))
+                {
+                    sb.AppendLine($"{Indent}{Indent}if (({variantListString}) -contains $parameterSet -and -not $PSBoundParameters.ContainsKey('{parameterName}'){setCondition}) {{");
+                    sb.AppendLine($"{Indent}{Indent}{Indent}$testPlayback = $false");
+                    sb.AppendLine($"{Indent}{Indent}{Indent}$PSBoundParameters['HttpPipelinePrepend'] | Foreach-Object {{ $testPlayback = $testPlayback -or ('Microsoft.Message.ClientRuntime.PipelineMock' -eq $_.Target.GetType().FullName) }}");
+                    sb.AppendLine($"{Indent}{Indent}{Indent}if ($testPlayback) {{");
+                    sb.AppendLine($"{Indent}{Indent}{Indent}{Indent}$PSBoundParameters['{parameterName}'] = . (Join-Path $PSScriptRoot '..' 'utils' 'Get-SubscriptionIdTestSafe.ps1')");
+                    sb.AppendLine($"{Indent}{Indent}{Indent}}} else {{");
+                    sb.AppendLine($"{Indent}{Indent}{Indent}{Indent}$PSBoundParameters['{parameterName}'] = {defaultInfo.Script}");
+                    sb.AppendLine($"{Indent}{Indent}{Indent}}}");
+                    sb.Append($"{Indent}{Indent}}}");
+                }
+                else
+                {
+                    sb.AppendLine($"{Indent}{Indent}if (({variantListString}) -contains $parameterSet -and -not $PSBoundParameters.ContainsKey('{parameterName}'){setCondition}) {{");
+                    sb.AppendLine($"{Indent}{Indent}{Indent}$PSBoundParameters['{parameterName}'] = {defaultInfo.Script}");
+                    sb.Append($"{Indent}{Indent}}}");
+                }
+
             }
             return sb.ToString();
         }

diff --git a/powershell/resources/utils/Get-SubscriptionIdTestSafe.ps1 b/powershell/resources/utils/Get-SubscriptionIdTestSafe.ps1
@@ -0,0 +1,7 @@
+param()
+if ($env:AzPSAutorestTestPlaybackMode) {
+    $loadEnvPath = Join-Path $PSScriptRoot '..' 'test' 'loadEnv.ps1'
+    . ($loadEnvPath)
+    return $env.SubscriptionId
+}
+return (Get-AzContext).Subscription.Id