refactoring HeaderMiddleware tests

Author: brettsam

src/WebJobs.Script.WebHost/Middleware/CustomHttpHeadersMiddleware.cs

@@ -21,12 +21,17 @@ public CustomHttpHeadersMiddleware(IOptions<CustomHttpHeadersOptions> hostOption
 
         public async Task Invoke(HttpContext context, RequestDelegate next)
         {
-            await next(context);
-
-            foreach (var header in _hostOptions)
+            context.Response.OnStarting(() =>
             {
-                context.Response.Headers.TryAdd(header.Key, header.Value);
-            }
+                foreach (var header in _hostOptions)
+                {
+                    context.Response.Headers.TryAdd(header.Key, header.Value);
+                }
+
+                return Task.CompletedTask;
+            });
+
+            await next(context);
         }
     }
-}
+}

test/WebJobs.Script.Tests/Middleware/CustomHttpHeadersMiddlewareTests.cs

@@ -3,19 +3,30 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Net;
+using System.Net.Http;
 using System.Text;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Features;
+using Microsoft.AspNetCore.TestHost;
+using Microsoft.Azure.WebJobs.Script.Middleware;
 using Microsoft.Azure.WebJobs.Script.WebHost.Configuration;
 using Microsoft.Azure.WebJobs.Script.WebHost.Middleware;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Options;
 using Xunit;
 
 namespace Microsoft.Azure.WebJobs.Script.Tests.Middleware
 {
     public class CustomHttpHeadersMiddlewareTests
     {
+        private bool _nextInvoked = false;
+
         [Fact]
         public async Task Invoke_hasCustomHeaders_AddsResponseHeaders()
         {
@@ -26,42 +37,52 @@ public async Task Invoke_hasCustomHeaders_AddsResponseHeaders()
             };
             var headerOptions = new OptionsWrapper<CustomHttpHeadersOptions>(headers);
 
-            bool nextInvoked = false;
-            RequestDelegate next = (context) =>
+            HttpClient client = GetTestHttpClient(o =>
             {
-                nextInvoked = true;
-                context.Response.StatusCode = (int)HttpStatusCode.Accepted;
-                return Task.CompletedTask;
-            };
+                o.Add("X-Content-Type-Options", "nosniff");
+                o.Add("Feature-Policy", "camera 'none'; geolocation 'none'");
+            });
+            HttpResponseMessage response = await client.GetAsync(string.Empty);
 
-            var middleware = new CustomHttpHeadersMiddleware(headerOptions);
-
-            var httpContext = new DefaultHttpContext();
-            await middleware.Invoke(httpContext, next);
-            Assert.True(nextInvoked);
-            Assert.Equal(httpContext.Response.Headers["X-Content-Type-Options"].ToString(), "nosniff");
-            Assert.Equal(httpContext.Response.Headers["Feature-Policy"].ToString(), "camera 'none'; geolocation 'none'");
+            Assert.True(_nextInvoked);
+            Assert.Equal(response.Headers.GetValues("X-Content-Type-Options").Single(), "nosniff");
+            Assert.Equal(response.Headers.GetValues("Feature-Policy").Single(), "camera 'none'; geolocation 'none'");
         }
 
         [Fact]
         public async Task Invoke_noCustomHeaders_DoesNotAddResponseHeader()
         {
-            var headerOptions = new OptionsWrapper<CustomHttpHeadersOptions>(new CustomHttpHeadersOptions());
-
-            bool nextInvoked = false;
-            RequestDelegate next = (context) =>
-            {
-                nextInvoked = true;
-                context.Response.StatusCode = (int)HttpStatusCode.Accepted;
-                return Task.CompletedTask;
-            };
+            HttpResponseMessage response = await GetTestHttpClient().GetAsync(string.Empty);
+            Assert.True(_nextInvoked);
+            Assert.Equal(response.Headers.Count(), 0);
+        }
 
-            var middleware = new CustomHttpHeadersMiddleware(headerOptions);
+        private HttpClient GetTestHttpClient(Action<CustomHttpHeadersOptions> configureOptions = null)
+        {
+            // The custom middleware relies on the host starting the request (thus invoking OnStarting),
+            // so we need to create a test host to flow through the entire pipeline.
+            var host = new WebHostBuilder()
+                .UseTestServer()
+                .ConfigureServices(s =>
+                {
+                    s.AddSingleton<IJobHostMiddlewarePipeline, DefaultMiddlewarePipeline>();
+                    s.AddSingleton<IJobHostHttpMiddleware, CustomHttpHeadersMiddleware>();
+                    s.AddOptions<CustomHttpHeadersOptions>().Configure(o => configureOptions?.Invoke(o));
+                })
+                .Configure(app =>
+                {
+                    app.UseMiddleware<JobHostPipelineMiddleware>();
+                    app.Use((context, next) =>
+                    {
+                        _nextInvoked = true;
+                        context.Response.StatusCode = (int)HttpStatusCode.Accepted;
+                        return Task.CompletedTask;
+                    });
+                })
+                .Build();
 
-            var httpContext = new DefaultHttpContext();
-            await middleware.Invoke(httpContext, next);
-            Assert.True(nextInvoked);
-            Assert.Equal(httpContext.Response.Headers.Count, 0);
+            host.Start();
+            return host.GetTestClient();
         }
     }
 }