From 097db04ae535b085945d7b5d58ec62eacad048f6 Mon Sep 17 00:00:00 2001
From: griffri <Griffri.Melville@digirati.com>
Date: Thu, 1 Aug 2024 17:35:26 +0100
Subject: [PATCH] Update PathRewriteTransformer to use IsSuccessStatusCode() to
 determine if a shorter `max-age` should be added to the header, update
 associated tests

---
 .../ReverseProxy/PathRewriteTransformerTests.cs        |  2 ++
 .../ReverseProxy/PathRewriteTransformer.cs             | 10 ++--------
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/src/protagonist/Orchestrator.Tests/Infrastructure/ReverseProxy/PathRewriteTransformerTests.cs b/src/protagonist/Orchestrator.Tests/Infrastructure/ReverseProxy/PathRewriteTransformerTests.cs
index c11181643..880310c22 100644
--- a/src/protagonist/Orchestrator.Tests/Infrastructure/ReverseProxy/PathRewriteTransformerTests.cs
+++ b/src/protagonist/Orchestrator.Tests/Infrastructure/ReverseProxy/PathRewriteTransformerTests.cs
@@ -166,6 +166,8 @@ public async Task TransformResponseAsync_AddsCacheHeaders_IfImageServer_AndPubli
     [InlineData(HttpStatusCode.BadGateway)]
     [InlineData(HttpStatusCode.ServiceUnavailable)]
     [InlineData(HttpStatusCode.GatewayTimeout)]
+    [InlineData(HttpStatusCode.NotFound)]
+    [InlineData(HttpStatusCode.BadRequest)]
     public async Task TransformResponseAsync_AddsSmallMaxAge_RegardlessOfDestination_IfError(HttpStatusCode statusCode)
     {
         // Arrange
diff --git a/src/protagonist/Orchestrator/Infrastructure/ReverseProxy/PathRewriteTransformer.cs b/src/protagonist/Orchestrator/Infrastructure/ReverseProxy/PathRewriteTransformer.cs
index 3dfbbcd38..bd84f161d 100644
--- a/src/protagonist/Orchestrator/Infrastructure/ReverseProxy/PathRewriteTransformer.cs
+++ b/src/protagonist/Orchestrator/Infrastructure/ReverseProxy/PathRewriteTransformer.cs
@@ -61,14 +61,8 @@ public override ValueTask<bool> TransformResponseAsync(
         return new ValueTask<bool>(true);
     }
 
-    private bool IsDownstreamError(HttpResponseMessage? proxyResponse)
-    {
-        var downstreamStatus = proxyResponse?.StatusCode ?? HttpStatusCode.InternalServerError;
-        return downstreamStatus is HttpStatusCode.InternalServerError
-            or HttpStatusCode.BadGateway
-            or HttpStatusCode.ServiceUnavailable
-            or HttpStatusCode.GatewayTimeout;
-    }
+    private bool IsDownstreamError(HttpResponseMessage? proxyResponse) 
+        => proxyResponse == null || !proxyResponse.IsSuccessStatusCode;
 
     private void EnsureCacheHeaders(HttpContext httpContext, bool isDownstreamError)
     {