From 7e69a4df73f57907a9439d95575b9458e1fa2b14 Mon Sep 17 00:00:00 2001
From: makharch <makharch@microsoft.com>
Date: Thu, 27 May 2021 18:28:08 -0700
Subject: [PATCH] Added output type fix and updated tests Changes from PR
 https://github.com/Azure/azure-powershell/pull/15125

---
 .../ScenarioTests/ResourceGraphQueryTests.ps1 | 18 +++++-----
 .../ResourceGraph/Az.ResourceGraph.psd1       |  2 +-
 .../Cmdlets/SearchAzureRmGraph.cs             |  4 +--
 .../Models/PSResourceGraphResponse.cs         | 36 +++++++++++++++++--
 4 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/src/ResourceGraph/ResourceGraph.Test/ScenarioTests/ResourceGraphQueryTests.ps1 b/src/ResourceGraph/ResourceGraph.Test/ScenarioTests/ResourceGraphQueryTests.ps1
index baa48a3dfb79..71cdff1996e4 100644
--- 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
index 29b3d7e14b69..f033d314e215 100644
--- 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
index 19ef4762a0d1..a8d7a79313ca 100644
--- 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
index 96e81b15eb35..66f75638bde9 100644
--- 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);
     }
 }