Added output type fix and updated tests

Changes from PR Azure#15125
fenwickt · Jun 9, 2021 · 98014e0 · 98014e0
1 parent cc484e5
commit 98014e0
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 14 deletions.
diff --git a/src/ResourceGraph/ResourceGraph.Test/ScenarioTests/ResourceGraphQueryTests.ps1 b/src/ResourceGraph/ResourceGraph.Test/ScenarioTests/ResourceGraphQueryTests.ps1
@@ -19,8 +19,8 @@ Run simple query
 function Search-AzureRmGraph-Query
 {
 	$queryResult = Search-AzGraph 'Resources | where tags != "" | project id, tags, properties | limit 2'
-
-	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse $queryResult
+	
+	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse[PSObject] $queryResult
 	Assert-Null $queryResult.SkipToken
 	Assert-IsInstance System.Collections.Generic.List[PSObject] $queryResult.Data
 	Assert-AreEqual 2 $queryResult.Data.Count
@@ -55,7 +55,7 @@ function Search-AzureRmGraph-PagedQuery
 	# Page size was artificially set to 2 rows
 	$queryResult = Search-AzGraph "project id" -First 3 -Skip 2
 
-	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse $queryResult
+	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse[PSObject] $queryResult
 	Assert-IsInstance System.String $queryResult.SkipToken
 	Assert-IsInstance System.Collections.Generic.List[PSObject] $queryResult.Data
 	Assert-AreEqual 3 $queryResult.Data.Count
@@ -95,17 +95,17 @@ function Search-AzureRmGraph-Subscriptions
 	$queryResultOneSub = Search-AzGraph $query -Subscription $testSubId
 	$queryResultMultipleSubs = Search-AzGraph $query -Subscription @($testSubId, $nonExsitentTestSubId)
 
-	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse $queryResultSubsFromContext
+	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse[PSObject] $queryResultSubsFromContext
 	Assert-Null $queryResultSubsFromContext.SkipToken
 	Assert-IsInstance System.Collections.Generic.List[PSObject] $queryResultSubsFromContext.Data
 	Assert-AreEqual $testSubId $queryResultSubsFromContext.Data.subscriptionId
 
-	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse $queryResultOneSub
+	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse[PSObject] $queryResultOneSub
 	Assert-Null $queryResultOneSub.SkipToken
 	Assert-IsInstance System.Collections.Generic.List[PSObject] $queryResultOneSub.Data
 	Assert-AreEqual $testSubId $queryResultOneSub.Data.subscriptionId
 
-	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse $queryResultMultipleSubs
+	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse[PSObject] $queryResultMultipleSubs
 	Assert-Null $queryResultMultipleSubs.SkipToken
 	Assert-IsInstance System.Collections.Generic.List[PSObject] $queryResultMultipleSubs.Data
 	Assert-AreEqual $testSubId $queryResultMultipleSubs.Data.subscriptionId
@@ -126,12 +126,12 @@ function Search-AzureRmGraph-ManagementGroups
 	$queryResultOneMg = Search-AzGraph $query -ManagementGroup $testMgId1
 	$queryResultMultipleMgs = Search-AzGraph $query -ManagementGroup @($testMgId1, $testMgId2, $nonExistentTestMgId) -AllowPartialScope
 
-	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse $queryResultOneMg	
+	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse[PSObject] $queryResultOneMg	
 	Assert-Null $queryResultOneMg.SkipToken
 	Assert-IsInstance System.Collections.Generic.List[PSObject] $queryResultOneMg.Data
 	Assert-AreEqual $testSubId $queryResultOneMg.Data.subscriptionId
 
-	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse $queryResultMultipleMgs
+	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse[PSObject] $queryResultMultipleMgs
 	Assert-Null $queryResultMultipleMgs.SkipToken
 	Assert-IsInstance System.Collections.Generic.List[PSObject] $queryResultMultipleMgs.Data
 	Assert-AreEqual $testSubId $queryResultMultipleMgs.Data.subscriptionId
@@ -145,7 +145,7 @@ function Search-AzureRmGraph-SkipTokenQuery
 {
 	$queryResult = Search-AzGraph "project id, properties" -SkipToken "ew0KICAiJGlkIjogIjEiLA0KICAiTWF4Um93cyI6IDMsDQogICJSb3dzVG9Ta2lwIjogMywNCiAgIkt1c3RvQ2x1c3RlclVybCI6ICJodHRwczovL2FyZy1ldXMtc2l4LXNmLmFyZy5jb3JlLndpbmRvd3MubmV0Ig0KfQ=="
 
-	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse $queryResult
+	Assert-IsInstance Microsoft.Azure.Commands.ResourceGraph.Models.PSResourceGraphResponse[PSObject] $queryResult
 	Assert-IsInstance System.String $queryResult.SkipToken
 	Assert-IsInstance System.Collections.Generic.List[PSObject] $queryResult.Data
 	Assert-AreEqual 3 $queryResult.Data.Count

diff --git a/src/ResourceGraph/ResourceGraph/Az.ResourceGraph.psd1 b/src/ResourceGraph/ResourceGraph/Az.ResourceGraph.psd1
@@ -70,7 +70,7 @@ FormatsToProcess = 'ResourceGraph.format.ps1xml',
                'ResourceGraph.Autorest\Az.ResourceGraph.format.ps1xml'
 
 # Modules to import as nested modules of the module specified in RootModule/ModuleToProcess
-NestedModules = @('Az.ResourceGraph.psm1', 
+NestedModules = @('Microsoft.Azure.PowerShell.Cmdlets.ResourceGraph.dll', 
                'ResourceGraph.Autorest\Az.ResourceGraph.psm1')
 
 # Functions to export from this module, for best performance, do not use wildcards and do not delete the entry, use an empty array if there are no functions to export.

diff --git a/src/ResourceGraph/ResourceGraph/Cmdlets/SearchAzureRmGraph.cs b/src/ResourceGraph/ResourceGraph/Cmdlets/SearchAzureRmGraph.cs
@@ -27,7 +27,7 @@ namespace Microsoft.Azure.Commands.ResourceGraph.Cmdlets
     /// Search-AzGraph cmdlet
     /// </summary>
     /// <seealso cref="Microsoft.Azure.Commands.ResourceGraph.Utilities.ResourceGraphBaseCmdlet" />
-    [Cmdlet(VerbsCommon.Search, ResourceManager.Common.AzureRMConstants.AzureRMPrefix + "Graph", DefaultParameterSetName = "SubscriptionScopedQuery"), OutputType(typeof(PSResourceGraphResponse))]
+    [Cmdlet(VerbsCommon.Search, ResourceManager.Common.AzureRMConstants.AzureRMPrefix + "Graph", DefaultParameterSetName = "SubscriptionScopedQuery"), OutputType(typeof(PSResourceGraphResponse<PSObject>))]
     public class SearchAzureRmGraph : ResourceGraphBaseCmdlet
     {
         /// <summary>
@@ -160,7 +160,7 @@ public override void ExecuteCmdlet()
                 }
             }
 
-            var psResourceGraphResponse = new PSResourceGraphResponse();
+            var psResourceGraphResponse = new PSResourceGraphResponse<PSObject>();
             QueryResponse response = null;
 
             var resultTruncated = false;

diff --git a/src/ResourceGraph/ResourceGraph/Models/PSResourceGraphResponse.cs b/src/ResourceGraph/ResourceGraph/Models/PSResourceGraphResponse.cs
@@ -14,17 +14,49 @@
 
 namespace Microsoft.Azure.Commands.ResourceGraph.Models
 {
+    using System.Collections;
     using System.Collections.Generic;
-    using System.Management.Automation;
     using Microsoft.WindowsAzure.Commands.Common.Attributes;
 
-    public class PSResourceGraphResponse
+    public class PSResourceGraphResponse<PSObject> : IList<PSObject>
     {
         [Ps1Xml(Target = ViewControl.List)]
         public string SkipToken { get; set; }
 
         [Ps1Xml(Target = ViewControl.List)]
         public IList<PSObject> Data { get; set; }
+        public PSObject this[int index] { get => Data[index]; set => Data[index] = value; }
 
+        public IEnumerator<PSObject> GetEnumerator()
+        {
+            return Data.GetEnumerator();
+        }
+
+        IEnumerator IEnumerable.GetEnumerator()
+        {
+            return GetEnumerator();
+        }
+
+        public bool IsReadOnly => Data.IsReadOnly;
+
+        public int Count => Data.Count;
+
+        public void Add(PSObject value) => Data.Add(value);
+
+        public void Clear() => Data.Clear();
+
+        public bool Contains(PSObject value) => Data.Contains(value);
+
+        public void CopyTo(PSObject[] array, int index) => Data.CopyTo(array, index);
+
+        public int IndexOf(PSObject value) => Data.IndexOf(value);
+
+        public void Insert(int index, PSObject value) => Data.Insert(index, value);
+
+        public void Remove(PSObject value) => Data.Remove(value);
+
+        public void RemoveAt(int index) => Data.RemoveAt(index);
+
+        bool ICollection<PSObject>.Remove(PSObject item) => Data.Remove(item);
     }
 }