Skip to content
42 changes: 20 additions & 22 deletions src/Middleware/ResponseCaching/test/CachedResponseBodyTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,13 @@
// The .NET Foundation licenses this file to you under the MIT license.

using System.Buffers;
using System.Diagnostics;
using System.IO.Pipelines;
using Microsoft.AspNetCore.InternalTesting;

namespace Microsoft.AspNetCore.ResponseCaching.Tests;

public class CachedResponseBodyTests
{
private readonly int _timeout = Debugger.IsAttached ? -1 : 5000;

[Fact]
public void GetSegments()
{
Expand All @@ -38,12 +35,11 @@ public async Task Copy_DoNothingWhenNoSegments()
var body = new CachedResponseBody(segments, 0);

var pipe = new Pipe();
using var cts = new CancellationTokenSource(_timeout);

var receiverTask = ReceiveDataAsync(pipe.Reader, receivedSegments, cts.Token);
var copyTask = body.CopyToAsync(pipe.Writer, cts.Token).ContinueWith(_ => pipe.Writer.CompleteAsync());
var receiverTask = ReceiveDataAsync(pipe.Reader, receivedSegments);
var copyTask = CopyDataAsync(body, pipe.Writer);

await Task.WhenAll(receiverTask, copyTask);
await Task.WhenAll(receiverTask, copyTask).DefaultTimeout();

Assert.Empty(receivedSegments);
}
Expand All @@ -60,12 +56,10 @@ public async Task Copy_SingleSegment()

var pipe = new Pipe();

using var cts = new CancellationTokenSource(_timeout);

var receiverTask = ReceiveDataAsync(pipe.Reader, receivedSegments, cts.Token);
var copyTask = CopyDataAsync(body, pipe.Writer, cts.Token);
var receiverTask = ReceiveDataAsync(pipe.Reader, receivedSegments);
var copyTask = CopyDataAsync(body, pipe.Writer);

await Task.WhenAll(receiverTask, copyTask);
await Task.WhenAll(receiverTask, copyTask).DefaultTimeout();

Assert.Equal(segments, receivedSegments);
}
Expand All @@ -83,27 +77,31 @@ public async Task Copy_MultipleSegments()

var pipe = new Pipe();

using var cts = new CancellationTokenSource(_timeout);
var receiverTask = ReceiveDataAsync(pipe.Reader, receivedSegments);
var copyTask = CopyDataAsync(body, pipe.Writer);

var receiverTask = ReceiveDataAsync(pipe.Reader, receivedSegments, cts.Token);
var copyTask = CopyDataAsync(body, pipe.Writer, cts.Token);

await Task.WhenAll(receiverTask, copyTask);
await Task.WhenAll(receiverTask, copyTask).DefaultTimeout();

Assert.Equal(new byte[] { 1, 2, 3 }, receivedSegments.SelectMany(x => x).ToArray());
}

static async Task CopyDataAsync(CachedResponseBody body, PipeWriter writer, CancellationToken cancellationToken)
static async Task CopyDataAsync(CachedResponseBody body, PipeWriter writer)
{
await body.CopyToAsync(writer, cancellationToken);
await writer.CompleteAsync();
try
{
await body.CopyToAsync(writer, CancellationToken.None);
}
finally
{
await writer.CompleteAsync();
}
}

static async Task ReceiveDataAsync(PipeReader reader, List<byte[]> receivedSegments, CancellationToken cancellationToken)
static async Task ReceiveDataAsync(PipeReader reader, List<byte[]> receivedSegments)
{
while (true)
{
var result = await reader.ReadAsync(cancellationToken);
var result = await reader.ReadAsync();
var buffer = result.Buffer;

foreach (var memory in buffer)
Expand Down
Loading