Use AsyncCollection in shared helper (#7055)

sdk/appconfiguration/Azure.ApplicationModel.Configuration/src/Azure.ApplicationModel.Configuration.csproj

@@ -26,7 +26,6 @@
   <ItemGroup>
     <Compile Include="$(AzureCoreSharedSources)ArrayBufferWriter.cs" />
     <Compile Include="$(AzureCoreSharedSources)HashCodeBuilder.cs" />
-    <Compile Include="$(AzureCoreSharedSources)PageResponse.cs" />
     <Compile Include="$(AzureCoreSharedSources)PageResponseEnumerator.cs" />
   </ItemGroup>
 

sdk/appconfiguration/Azure.ApplicationModel.Configuration/src/ConfigurationClient.cs

@@ -569,7 +569,7 @@ public virtual Response<ConfigurationSetting> Get(string key, string label = def
         /// </summary>
         /// <param name="selector">Set of options for selecting <see cref="ConfigurationSetting"/> from the configuration store.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
-        public virtual IAsyncEnumerable<Response<ConfigurationSetting>> GetSettingsAsync(SettingSelector selector, CancellationToken cancellationToken = default)
+        public virtual AsyncCollection<ConfigurationSetting> GetSettingsAsync(SettingSelector selector, CancellationToken cancellationToken = default)
         {
             return PageResponseEnumerator.CreateAsyncEnumerable(nextLink => GetSettingsPageAsync(selector, nextLink, cancellationToken));
         }
@@ -589,7 +589,7 @@ public virtual IEnumerable<Response<ConfigurationSetting>> GetSettings(SettingSe
         /// </summary>
         /// <param name="selector">Set of options for selecting <see cref="ConfigurationSetting"/> from the configuration store.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
-        public virtual IAsyncEnumerable<Response<ConfigurationSetting>> GetRevisionsAsync(SettingSelector selector, CancellationToken cancellationToken = default)
+        public virtual AsyncCollection<ConfigurationSetting> GetRevisionsAsync(SettingSelector selector, CancellationToken cancellationToken = default)
         {
             return PageResponseEnumerator.CreateAsyncEnumerable(nextLink => GetRevisionsPageAsync(selector, nextLink, cancellationToken));
         }
@@ -630,7 +630,7 @@ private Request CreateGetRequest(string key, string label, DateTimeOffset accept
         /// <param name="selector">Set of options for selecting settings from the configuration store.</param>
         /// <param name="pageLink"></param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
-        private async Task<PageResponse<ConfigurationSetting>> GetSettingsPageAsync(SettingSelector selector, string pageLink, CancellationToken cancellationToken = default)
+        private async Task<Page<ConfigurationSetting>> GetSettingsPageAsync(SettingSelector selector, string pageLink, CancellationToken cancellationToken = default)
         {
             using DiagnosticScope scope = _pipeline.Diagnostics.CreateScope("Azure.ApplicationModel.Configuration.ConfigurationClient.GetSettingsPage");
             scope.Start();
@@ -645,7 +645,7 @@ private async Task<PageResponse<ConfigurationSetting>> GetSettingsPageAsync(Sett
                     case 200:
                     case 206:
                         SettingBatch settingBatch = await ConfigurationServiceSerializer.ParseBatchAsync(response, cancellationToken).ConfigureAwait(false);
-                        return new PageResponse<ConfigurationSetting>(settingBatch.Settings, response, settingBatch.NextBatchLink);
+                        return new Page<ConfigurationSetting>(settingBatch.Settings, settingBatch.NextBatchLink, response);
                     default:
                         throw await response.CreateRequestFailedExceptionAsync().ConfigureAwait(false);
                 }
@@ -663,7 +663,7 @@ private async Task<PageResponse<ConfigurationSetting>> GetSettingsPageAsync(Sett
         /// <param name="selector">Set of options for selecting settings from the configuration store.</param>
         /// <param name="pageLink"></param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
-        private PageResponse<ConfigurationSetting> GetSettingsPage(SettingSelector selector, string pageLink, CancellationToken cancellationToken = default)
+        private Page<ConfigurationSetting> GetSettingsPage(SettingSelector selector, string pageLink, CancellationToken cancellationToken = default)
         {
             using DiagnosticScope scope = _pipeline.Diagnostics.CreateScope("Azure.ApplicationModel.Configuration.ConfigurationClient.GetSettingsPage");
             scope.Start();
@@ -678,7 +678,7 @@ private PageResponse<ConfigurationSetting> GetSettingsPage(SettingSelector selec
                     case 200:
                     case 206:
                         SettingBatch settingBatch = ConfigurationServiceSerializer.ParseBatch(response);
-                        return new PageResponse<ConfigurationSetting>(settingBatch.Settings, response, settingBatch.NextBatchLink);
+                        return new Page<ConfigurationSetting>(settingBatch.Settings, settingBatch.NextBatchLink, response);
                     default:
                         throw response.CreateRequestFailedException();
                 }
@@ -712,7 +712,7 @@ private Request CreateBatchRequest(SettingSelector selector, string pageLink)
         /// <param name="selector">Set of options for selecting settings from the configuration store.</param>
         /// <param name="pageLink"></param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
-        private async Task<PageResponse<ConfigurationSetting>> GetRevisionsPageAsync(SettingSelector selector, string pageLink, CancellationToken cancellationToken = default)
+        private async Task<Page<ConfigurationSetting>> GetRevisionsPageAsync(SettingSelector selector, string pageLink, CancellationToken cancellationToken = default)
         {
             using DiagnosticScope scope = _pipeline.Diagnostics.CreateScope("Azure.ApplicationModel.Configuration.ConfigurationClient.GetRevisionsPage");
             scope.Start();
@@ -726,7 +726,7 @@ private async Task<PageResponse<ConfigurationSetting>> GetRevisionsPageAsync(Set
                     case 200:
                     case 206:
                         SettingBatch settingBatch = await ConfigurationServiceSerializer.ParseBatchAsync(response, cancellationToken).ConfigureAwait(false);
-                        return new PageResponse<ConfigurationSetting>(settingBatch.Settings, response, settingBatch.NextBatchLink);
+                        return new Page<ConfigurationSetting>(settingBatch.Settings, settingBatch.NextBatchLink, response);
                     default:
                         throw await response.CreateRequestFailedExceptionAsync().ConfigureAwait(false);
                 }
@@ -745,7 +745,7 @@ private async Task<PageResponse<ConfigurationSetting>> GetRevisionsPageAsync(Set
         /// <param name="selector">Set of options for selecting settings from the configuration store.</param>
         /// <param name="pageLink"></param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
-        private PageResponse<ConfigurationSetting> GetRevisionsPage(SettingSelector selector, string pageLink, CancellationToken cancellationToken = default)
+        private Page<ConfigurationSetting> GetRevisionsPage(SettingSelector selector, string pageLink, CancellationToken cancellationToken = default)
         {
             using DiagnosticScope scope = _pipeline.Diagnostics.CreateScope("Azure.ApplicationModel.Configuration.ConfigurationClient.GetRevisionsPage");
             scope.Start();
@@ -759,7 +759,7 @@ private PageResponse<ConfigurationSetting> GetRevisionsPage(SettingSelector sele
                     case 200:
                     case 206:
                         SettingBatch settingBatch = ConfigurationServiceSerializer.ParseBatch(response);
-                        return new PageResponse<ConfigurationSetting>(settingBatch.Settings, response, settingBatch.NextBatchLink);
+                        return new Page<ConfigurationSetting>(settingBatch.Settings, settingBatch.NextBatchLink, response);
                     default:
                         throw response.CreateRequestFailedException();
                 }

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

@@ -4,6 +4,7 @@
 using System.Collections.Generic;
 using System.ComponentModel;
 using System.Threading;
+using System.Threading.Tasks;
 
 namespace Azure
 {
@@ -66,7 +67,16 @@ public abstract IAsyncEnumerable<Page<T>> ByPage(
         /// enumerating asynchronously.
         /// </param>
         /// <returns>An async sequence of values.</returns>
-        public abstract IAsyncEnumerator<Response<T>> GetAsyncEnumerator(CancellationToken cancellationToken = default);
+        public virtual async IAsyncEnumerator<Response<T>> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+        {
+            await foreach (Page<T> page in ByPage().ConfigureAwait(false).WithCancellation(cancellationToken))
+            {
+                foreach (T value in page.Values)
+                {
+                    yield return new Response<T>(page.GetRawResponse(), value);
+                }
+            }
+        }
 
         /// <summary>
         /// Creates a string representation of an <see cref="AsyncCollection{T}"/>.

sdk/core/Azure.Core/src/Shared/PageResponseEnumerator.cs

@@ -2,39 +2,57 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections;
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Azure.Core
 {
     internal static class PageResponseEnumerator
     {
-        public static IEnumerable<Response<T>> CreateEnumerable<T>(Func<string, PageResponse<T>> pageFunc)
+        public static IEnumerable<Response<T>> CreateEnumerable<T>(Func<string, Page<T>> pageFunc)
         {
             string nextLink = null;
             do
             {
-                PageResponse<T> pageResponse = pageFunc(nextLink);
+                Page<T> pageResponse = pageFunc(nextLink);
                 foreach (T setting in pageResponse.Values)
                 {
-                    yield return new Response<T>(pageResponse.Response, setting);
+                    yield return new Response<T>(pageResponse.GetRawResponse(), setting);
                 }
-                nextLink = pageResponse.NextLink;
+                nextLink = pageResponse.ContinuationToken;
             } while (nextLink != null);
         }
 
-        public static async IAsyncEnumerable<Response<T>> CreateAsyncEnumerable<T>(Func<string, Task<PageResponse<T>>> pageFunc)
+        public static AsyncCollection<T> CreateAsyncEnumerable<T>(Func<string, Task<Page<T>>> pageFunc)
         {
-            string nextLink = null;
-            do
+            return new FuncAsyncCollection<T>((continuationToken, pageSizeHint) => pageFunc(continuationToken));
+        }
+
+        public static AsyncCollection<T> CreateAsyncEnumerable<T>(Func<string, int?, Task<Page<T>>> pageFunc)
+        {
+            return new FuncAsyncCollection<T>(pageFunc);
+        }
+
+        internal class FuncAsyncCollection<T>: AsyncCollection<T>
+        {
+            private readonly Func<string, int?, Task<Page<T>>> _pageFunc;
+
+            public FuncAsyncCollection(Func<string, int?, Task<Page<T>>> pageFunc)
             {
-                PageResponse<T> pageResponse = await pageFunc(nextLink).ConfigureAwait(false);
-                foreach (T setting in pageResponse.Values)
+                _pageFunc = pageFunc;
+            }
+
+            public override async IAsyncEnumerable<Page<T>> ByPage(string continuationToken = default, int? pageSizeHint = default)
+            {
+                do
                 {
-                    yield return new Response<T>(pageResponse.Response, setting);
-                }
-                nextLink = pageResponse.NextLink;
-            } while (nextLink != null);
+                    Page<T> pageResponse = await _pageFunc(continuationToken, pageSizeHint).ConfigureAwait(false);
+                    yield return pageResponse;
+                    continuationToken = pageResponse.ContinuationToken;
+                } while (continuationToken != null);
+            }
         }
     }
 }

sdk/core/Azure.Core/tests/TestFramework/AsyncOnlyAttribute.cs

@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+using System.Text;
+using NUnit.Framework;
+using NUnit.Framework.Interfaces;
+using NUnit.Framework.Internal;
+
+namespace Azure.Core.Testing
+{
+
+    [AttributeUsage(AttributeTargets.Method|AttributeTargets.Class|AttributeTargets.Assembly, AllowMultiple=false, Inherited=false)]
+    public class AsyncOnlyAttribute : NUnitAttribute, IApplyToTest
+    {
+        #region IApplyToTest members
+
+        /// <summary>
+        /// Modifies a test by marking it as Ignored.
+        /// </summary>
+        /// <param name="test">The test to modify</param>
+        public void ApplyToTest(Test test)
+        {
+            if (test.RunState != RunState.NotRunnable)
+            {
+                // This is an unfortunate implementation but it's the only one I was able to figure out
+                string testParameters = test.FullName.Substring(test.ClassName.Length);
+                if (testParameters.StartsWith("(False"))
+                {
+                    test.RunState = RunState.Ignored;
+                }
+            }
+        }
+
+        #endregion
+    }
+}

sdk/core/Azure.Core/tests/TestFramework/UseSyncMethodsInterceptor.cs

@@ -82,20 +82,10 @@ public void Intercept(IInvocation invocation)
                 // Map IEnumerable to IAsyncEnumerable
                 if (returnsIEnumerable)
                 {
-                    if (invocation.Method.ReturnType.IsGenericType &&
-                        invocation.Method.ReturnType.GetGenericTypeDefinition().Name == "AsyncCollection`1")
-                    {
-                        // AsyncCollection can be used as either a sync or async
-                        // collection so there's no need to wrap it in an
-                        // IAsyncEnumerable
-                        invocation.ReturnValue = result;
-                    }
-                    else
-                    {
-                        invocation.ReturnValue = Activator.CreateInstance(
-                            typeof(AsyncEnumerableWrapper<>).MakeGenericType(returnType.GenericTypeArguments),
-                            new[] { result });
-                    }
+                    Type[] modelType = returnType.GenericTypeArguments.Single().GenericTypeArguments;
+                    Type wrapperType = typeof(AsyncEnumerableWrapper<>).MakeGenericType(modelType);
+
+                    invocation.ReturnValue = Activator.CreateInstance(wrapperType, new [] { result });
                 }
                 else
                 {
@@ -120,34 +110,52 @@ private static MethodInfo GetMethod(IInvocation invocation, string nonAsyncMetho
             return invocation.TargetType.GetMethod(nonAsyncMethodName, BindingFlags.Public | BindingFlags.Instance, null, types, null);
         }
 
-        private class AsyncEnumerableWrapper<T> : IAsyncEnumerable<T>
+        private class AsyncEnumerableWrapper<T> : AsyncCollection<T>
         {
-            private readonly IEnumerable<T> _enumerable;
+            private readonly IEnumerable<Response<T>> _enumerable;
 
-            public AsyncEnumerableWrapper(IEnumerable<T> enumerable)
+            public AsyncEnumerableWrapper(IEnumerable<Response<T>> enumerable)
             {
                 _enumerable = enumerable;
             }
 
-            public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = new CancellationToken())
+#pragma warning disable 1998
+            public override async IAsyncEnumerable<Page<T>> ByPage(string continuationToken = default, int? pageSizeHint = default)
+#pragma warning restore 1998
             {
-                return new Enumerator(_enumerable.GetEnumerator());
+                if (continuationToken != null)
+                {
+                    throw new InvalidOperationException("Calling ByPage with a continuationToken is not supported in the sync mode");
+                }
+
+                if (pageSizeHint != null)
+                {
+                    throw new InvalidOperationException("Calling ByPage with a pageSizeHint is not supported in the sync mode");
+                }
+
+                foreach (Response<T> response in _enumerable)
+                {
+                    yield return new Page<T>(new [] { response.Value}, null, response.GetRawResponse());
+                }
             }
 
-            private class Enumerator: IAsyncEnumerator<T>
+            private class SingleEnumerable: IAsyncEnumerable<Page<T>>, IAsyncEnumerator<Page<T>>
             {
-                private readonly IEnumerator<T> _enumerator;
-
-                public Enumerator(IEnumerator<T> enumerator)
+                public SingleEnumerable(Page<T> value)
                 {
-                    _enumerator = enumerator;
+                    Current = value;
                 }
 
                 public ValueTask DisposeAsync() => default;
 
-                public ValueTask<bool> MoveNextAsync() => new ValueTask<bool>(_enumerator.MoveNext());
+                public ValueTask<bool> MoveNextAsync() => new ValueTask<bool>(false);
 
-                public T Current => _enumerator.Current;
+                public Page<T> Current { get; }
+
+                public IAsyncEnumerator<Page<T>> GetAsyncEnumerator(CancellationToken cancellationToken = new CancellationToken())
+                {
+                    return this;
+                }
             }
         }
     }

sdk/keyvault/Azure.Security.KeyVault.Keys/src/Azure.Security.KeyVault.Keys.csproj

@@ -26,7 +26,6 @@
 
   <ItemGroup>
     <Compile Include="$(AzureCoreSharedSources)ArrayBufferWriter.cs" />
-    <Compile Include="$(AzureCoreSharedSources)PageResponse.cs" />
     <Compile Include="$(AzureCoreSharedSources)PageResponseEnumerator.cs" />
   </ItemGroup>