Add multi page support for Pageable.FromPages method Azure#22705 (Azure#22706)

Author: cyrildurand

sdk/core/Azure.Core/src/AsyncPageable.cs

@@ -145,7 +145,7 @@ public static AsyncPageable<T> FromPages(IEnumerable<Page<T>> pages)
         [EditorBrowsable(EditorBrowsableState.Never)]
         public override int GetHashCode() => base.GetHashCode();
 
-        private class StaticPageable: AsyncPageable<T>
+        private class StaticPageable : AsyncPageable<T>
         {
             private readonly IEnumerable<Page<T>> _pages;
 
@@ -158,9 +158,21 @@ public StaticPageable(IEnumerable<Page<T>> pages)
             public override async IAsyncEnumerable<Page<T>> AsPages(string? continuationToken = default, int? pageSizeHint = default)
 #pragma warning restore 1998
             {
+                var shouldReturnPages = continuationToken == null;
+
                 foreach (var page in _pages)
                 {
-                    yield return page;
+                    if (shouldReturnPages)
+                    {
+                        yield return page;
+                    }
+                    else
+                    {
+                        if (continuationToken == page.ContinuationToken)
+                        {
+                            shouldReturnPages = true;
+                        }
+                    }
                 }
             }
         }

sdk/core/Azure.Core/src/Pageable.cs

@@ -115,7 +115,7 @@ public static Pageable<T> FromPages(IEnumerable<Page<T>> pages)
         [EditorBrowsable(EditorBrowsableState.Never)]
         public override int GetHashCode() => base.GetHashCode();
 
-        private class StaticPageable: Pageable<T>
+        private class StaticPageable : Pageable<T>
         {
             private readonly IEnumerable<Page<T>> _pages;
 
@@ -126,7 +126,22 @@ public StaticPageable(IEnumerable<Page<T>> pages)
 
             public override IEnumerable<Page<T>> AsPages(string? continuationToken = default, int? pageSizeHint = default)
             {
-                return _pages;
+                var shouldReturnPages = continuationToken == null;
+
+                foreach (var page in _pages)
+                {
+                    if (shouldReturnPages)
+                    {
+                        yield return page;
+                    }
+                    else
+                    {
+                        if (continuationToken == page.ContinuationToken)
+                        {
+                            shouldReturnPages = true;
+                        }
+                    }
+                }
             }
         }
     }

sdk/core/Azure.Core/tests/PageableTests.cs

@@ -33,5 +33,49 @@ public async Task CanCreateAsyncPageableFromPages()
 
             Assert.AreEqual(new byte[] { 1, 2, 3, 4 }, await pageable.ToEnumerableAsync());
         }
+
+        [Test]
+        public void CanCreatePageableFromPage_MultiplePage()
+        {
+            var pageable = Pageable<byte>.FromPages(new[]
+            {
+                Page<byte>.FromValues(new byte[] {1, 2}, "X", null),
+                Page<byte>.FromValues(new byte[] {3, 4}, null, null)
+            });
+
+            var pages = pageable.AsPages(continuationToken: null);
+            var page = pages.First();
+
+            Assert.AreEqual(new byte[] { 1, 2 }, page.Values);
+            Assert.AreEqual("X", page.ContinuationToken);
+
+            pages = pageable.AsPages(continuationToken: "X");
+            page = pages.First();
+
+            Assert.AreEqual(new byte[] { 3, 4 }, page.Values);
+            Assert.Null(page.ContinuationToken);
+        }
+
+        [Test]
+        public async Task CanCreateAsyncPageableFromPage_MultiplePage()
+        {
+            var pageable = AsyncPageable<byte>.FromPages(new[]
+            {
+                Page<byte>.FromValues(new byte[] {1, 2}, "X", null),
+                Page<byte>.FromValues(new byte[] {3, 4}, null, null)
+            });
+
+            var pages = pageable.AsPages(continuationToken: null);
+            var page = await pages.FirstAsync();
+
+            Assert.AreEqual(new byte[] { 1, 2 }, page.Values);
+            Assert.AreEqual("X", page.ContinuationToken);
+
+            pages = pageable.AsPages(continuationToken: "X");
+            page = await pages.FirstAsync();
+
+            Assert.AreEqual(new byte[] { 3, 4 }, page.Values);
+            Assert.Null(page.ContinuationToken);
+        }
     }
-}
+}