From 2f3ac9b52111cbcd680d082c010b4f7e1414bc0b Mon Sep 17 00:00:00 2001
From: Dustin Campbell <dustin@teamcampbell.org>
Date: Thu, 21 Mar 2024 12:02:33 -0700
Subject: [PATCH] Add a few AsyncBatchingWorkQueue tests

---
 .../Utilities/AsyncBatchingWorkQueueTest.cs   | 111 ++++++++++++++++++
 1 file changed, 111 insertions(+)
 create mode 100644 src/Razor/test/Microsoft.CodeAnalysis.Razor.Workspaces.Test/Utilities/AsyncBatchingWorkQueueTest.cs

diff --git a/src/Razor/test/Microsoft.CodeAnalysis.Razor.Workspaces.Test/Utilities/AsyncBatchingWorkQueueTest.cs b/src/Razor/test/Microsoft.CodeAnalysis.Razor.Workspaces.Test/Utilities/AsyncBatchingWorkQueueTest.cs
new file mode 100644
index 00000000000..301e3cf2821
--- /dev/null
+++ b/src/Razor/test/Microsoft.CodeAnalysis.Razor.Workspaces.Test/Utilities/AsyncBatchingWorkQueueTest.cs
@@ -0,0 +1,111 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.Test.Common;
+using Microsoft.CodeAnalysis.Razor.Utilities;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace Microsoft.CodeAnalysis.Razor.Workspaces.Test.Utilities;
+
+public class AsyncBatchingWorkQueueTest(ITestOutputHelper output) : ToolingTestBase(output)
+{
+    [Fact]
+    public async Task AddItemsAndWaitForBatch()
+    {
+        var list = new List<int>();
+
+        var workQueue = new AsyncBatchingWorkQueue<int>(
+            delay: TimeSpan.FromMilliseconds(1),
+            processBatchAsync: (items, cancellationToken) =>
+            {
+                foreach (var item in items)
+                {
+                    list.Add(item);
+                }
+
+                return default;
+            },
+            DisposalToken);
+
+        for (var i = 0; i < 1000; i++)
+        {
+            workQueue.AddWork(i);
+        }
+
+        await workQueue.WaitUntilCurrentBatchCompletesAsync();
+
+        for (var i = 0; i < 1000; i++)
+        {
+            Assert.Equal(i, list[i]);
+        }
+    }
+
+    [Fact]
+    public async Task DedupesItems()
+    {
+        var list = new List<int>();
+
+        var workQueue = new AsyncBatchingWorkQueue<int>(
+            delay: TimeSpan.FromMilliseconds(1),
+            processBatchAsync: (items, cancellationToken) =>
+            {
+                foreach (var item in items)
+                {
+                    list.Add(item);
+                }
+
+                return default;
+            },
+            equalityComparer: EqualityComparer<int>.Default,
+            DisposalToken);
+
+        for (var i = 0; i < 1000; i++)
+        {
+            workQueue.AddWork(i);
+            workQueue.AddWork(i);
+        }
+
+        await workQueue.WaitUntilCurrentBatchCompletesAsync();
+
+        for (var i = 0; i < 1000; i++)
+        {
+            Assert.Equal(i, list[i]);
+        }
+    }
+
+    [Fact]
+    public async Task CancelExistingWork()
+    {
+        var cancelled = false;
+
+        var workQueue = new AsyncBatchingWorkQueue(
+            TimeSpan.FromMilliseconds(1),
+            cancellationToken =>
+            {
+                while (true) // infinite loop
+                {
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        cancelled = true;
+                        break;
+                    }
+                }
+
+                return default;
+            },
+            DisposalToken);
+
+        // Add first bit of work
+        workQueue.AddWork();
+        await Task.Delay(20);
+
+        // Add another bit of work, cancelling the previous work.
+        workQueue.AddWork(cancelExistingWork: true);
+
+        Assert.True(cancelled);
+    }
+}