From 18380e566e7ce09b39e329e7c7b02c55531c904e Mon Sep 17 00:00:00 2001
From: Marco Burato <marco@ecutek.com>
Date: Thu, 17 Aug 2023 09:22:40 +0200
Subject: [PATCH] Add new RestClient.DownloadStreamAsync() overload which
 accepts an error handler to process error responses

---
 src/RestSharp/RestClient.Async.cs             | 39 +++++++++++++--
 .../DownloadFileTests.cs                      | 49 +++++++++++++++++++
 2 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/src/RestSharp/RestClient.Async.cs b/src/RestSharp/RestClient.Async.cs
index f3b738293..760a2e4e8 100644
--- a/src/RestSharp/RestClient.Async.cs
+++ b/src/RestSharp/RestClient.Async.cs
@@ -40,19 +40,50 @@ public async Task<RestResponse> ExecuteAsync(RestRequest request, CancellationTo
     /// <inheritdoc />
     [PublicAPI]
     public async Task<Stream?> DownloadStreamAsync(RestRequest request, CancellationToken cancellationToken = default) {
+        return await DownloadStreamInternalAsync(request, null, cancellationToken);
+    }
+
+    /// <inheritdoc />
+    [PublicAPI]
+    public async Task<Stream?> DownloadStreamAsync(RestRequest request, Action<RestResponse> errorHandler, CancellationToken cancellationToken = default) {
+        return await DownloadStreamInternalAsync(request, errorHandler, cancellationToken);
+    }
+
+    async Task<Stream?> DownloadStreamInternalAsync(RestRequest request, Action<RestResponse>? errorHandler, CancellationToken cancellationToken = default) {
         // Make sure we only read the headers so we can stream the content body efficiently
         request.CompletionOption = HttpCompletionOption.ResponseHeadersRead;
-        var response = await ExecuteRequestAsync(request, cancellationToken).ConfigureAwait(false);
+        var internalResponse = await ExecuteRequestAsync(request, cancellationToken).ConfigureAwait(false);
+
+        if (errorHandler != null) {
+            if (internalResponse.Exception != null) {
+                var response = GetErrorResponse(request, internalResponse.Exception, internalResponse.TimeoutToken);
+
+                errorHandler(response);
+            }
+            else if (internalResponse.ResponseMessage!.IsSuccessStatusCode == false) {
+                var response = await RestResponse.FromHttpResponse(
+                    internalResponse.ResponseMessage!,
+                    request,
+                    Options.Encoding,
+                    internalResponse.CookieContainer?.GetCookies(internalResponse.Url),
+                    Options.CalculateResponseStatus,
+                    cancellationToken
+                )
+                .ConfigureAwait(false);
+
+                errorHandler(response);
+            }
+        }
 
-        var exception = response.Exception ?? response.ResponseMessage?.MaybeException();
+        var exception = internalResponse.Exception ?? internalResponse.ResponseMessage?.MaybeException();
 
         if (exception != null) {
             return Options.ThrowOnAnyError ? throw exception : null;
         }
 
-        if (response.ResponseMessage == null) return null;
+        if (internalResponse.ResponseMessage == null) return null;
 
-        return await response.ResponseMessage.ReadResponseStream(request.ResponseWriter, cancellationToken).ConfigureAwait(false);
+        return await internalResponse.ResponseMessage.ReadResponseStream(request.ResponseWriter, cancellationToken).ConfigureAwait(false);
     }
 
     static RestResponse GetErrorResponse(RestRequest request, Exception exception, CancellationToken timeoutToken) {
diff --git a/test/RestSharp.Tests.Integrated/DownloadFileTests.cs b/test/RestSharp.Tests.Integrated/DownloadFileTests.cs
index f2cd59482..5d0476427 100644
--- a/test/RestSharp.Tests.Integrated/DownloadFileTests.cs
+++ b/test/RestSharp.Tests.Integrated/DownloadFileTests.cs
@@ -1,5 +1,6 @@
 using System.Net;
 using System.Text;
+using RestSharp.Extensions;
 using RestSharp.Tests.Shared.Fixtures;
 
 namespace RestSharp.Tests.Integrated;
@@ -9,6 +10,7 @@ public DownloadFileTests() {
         _server = HttpServerFixture.StartServer("Assets/Koala.jpg", FileHandler);
         var options = new RestClientOptions(_server.Url) { ThrowOnAnyError = true };
         _client = new RestClient(options);
+        _clientNoThrow = new RestClient(_server.Url);
     }
 
     public void Dispose() => _server.Dispose();
@@ -28,6 +30,7 @@ void FileHandler(HttpListenerRequest request, HttpListenerResponse response) {
 
     readonly HttpServerFixture _server;
     readonly RestClient        _client;
+    readonly RestClient        _clientNoThrow;
     readonly string            _path = AppDomain.CurrentDomain.BaseDirectory;
 
     [Fact]
@@ -65,6 +68,52 @@ public async Task Handles_Binary_File_Download() {
         Assert.Equal(expected, response);
     }
 
+    [Fact]
+    public async Task Runs_ErrorHandler_On_Download_Request_Failure() {
+        var client = new RestClient("http://localhost:12345");
+        var request = new RestRequest("nonexisting");
+        RestResponse? errorResponse = null;
+        var stream = await client.DownloadStreamAsync(request, (r) => {
+            errorResponse = r;
+        });
+
+        Assert.Null(stream);
+        Assert.NotNull(errorResponse);
+        Assert.Equal(ResponseStatus.Error, errorResponse.ResponseStatus);
+    }
+
+    [Fact]
+    public async Task Runs_ErrorHandler_On_Download_Response_StatusCode_Not_Successful() {
+        var request = new RestRequest("Assets/Koala1.jpg");
+        RestResponse? errorResponse = null;
+        var stream = await _clientNoThrow.DownloadStreamAsync(request, (r) => {
+            errorResponse = r;
+        });
+
+        Assert.Null(stream);
+        Assert.NotNull(errorResponse);
+        Assert.Equal(ResponseStatus.Completed, errorResponse.ResponseStatus);
+        Assert.False(errorResponse.IsSuccessStatusCode);
+        Assert.Equal(HttpStatusCode.NotFound, errorResponse.StatusCode);
+    }
+
+    [Fact]
+    public async Task Doesnt_Run_ErrorHandler_On_Download_Success() {
+        var request = new RestRequest("Assets/Koala.jpg");
+        RestResponse? errorResponse = null;
+        var stream = await _clientNoThrow.DownloadStreamAsync(request, (r) => {
+            errorResponse = r;
+        });
+
+        Assert.NotNull(stream);
+        Assert.Null(errorResponse);
+
+        var expected = await File.ReadAllBytesAsync(Path.Combine(_path, "Assets", "Koala.jpg"));
+        var bytes = await stream.ReadAsBytes(CancellationToken.None);
+
+        Assert.Equal(expected, bytes);
+    }
+
     [Fact]
     public async Task Writes_Response_To_Stream() {
         var tempFile = Path.GetTempFileName();