Add CreateHandler method that allows additonal configuration (#42677)

Author: justindbaur

src/Hosting/TestHost/src/ClientHandler.cs

@@ -20,16 +20,19 @@ namespace Microsoft.AspNetCore.TestHost;
 public class ClientHandler : HttpMessageHandler
 {
     private readonly ApplicationWrapper _application;
+    private readonly Action<HttpContext> _additionalContextConfiguration;
     private readonly PathString _pathBase;
 
     /// <summary>
     /// Create a new handler.
     /// </summary>
     /// <param name="pathBase">The base path.</param>
     /// <param name="application">The <see cref="IHttpApplication{TContext}"/>.</param>
-    internal ClientHandler(PathString pathBase, ApplicationWrapper application)
+    /// <param name="additionalContextConfiguration">The action to additionally configure <see cref="HttpContext"/>.</param>
+    internal ClientHandler(PathString pathBase, ApplicationWrapper application, Action<HttpContext>? additionalContextConfiguration = null)
     {
         _application = application ?? throw new ArgumentNullException(nameof(application));
+        _additionalContextConfiguration = additionalContextConfiguration ?? NoExtraConfiguration;
 
         // PathString.StartsWithSegments that we use below requires the base path to not end in a slash.
         if (pathBase.HasValue && pathBase.Value.EndsWith('/'))
@@ -163,6 +166,8 @@ protected override async Task<HttpResponseMessage> SendAsync(
             req.QueryString = QueryString.FromUriComponent(request.RequestUri);
         });
 
+        contextBuilder.Configure((context, _) => _additionalContextConfiguration(context));
+
         var response = new HttpResponseMessage();
 
         // Copy trailers to the response message when the response stream is complete
@@ -200,4 +205,9 @@ protected override async Task<HttpResponseMessage> SendAsync(
         }
         return response;
     }
+
+    private static void NoExtraConfiguration(HttpContext context)
+    {
+        // Intentional no op
+    }
 }

src/Hosting/TestHost/src/PublicAPI.Unshipped.txt

@@ -1 +1,2 @@
 #nullable enable
+Microsoft.AspNetCore.TestHost.TestServer.CreateHandler(System.Action<Microsoft.AspNetCore.Http.HttpContext!>! additionalContextConfiguration) -> System.Net.Http.HttpMessageHandler!

src/Hosting/TestHost/src/TestServer.cs

@@ -146,13 +146,30 @@ private ApplicationWrapper Application
         get => _application ?? throw new InvalidOperationException("The server has not been started or no web application was configured.");
     }
 
+    private PathString PathBase => BaseAddress == null ? PathString.Empty : PathString.FromUriComponent(BaseAddress);
+
     /// <summary>
     /// Creates a custom <see cref="HttpMessageHandler" /> for processing HTTP requests/responses with the test server.
     /// </summary>
     public HttpMessageHandler CreateHandler()
     {
-        var pathBase = BaseAddress == null ? PathString.Empty : PathString.FromUriComponent(BaseAddress);
-        return new ClientHandler(pathBase, Application) { AllowSynchronousIO = AllowSynchronousIO, PreserveExecutionContext = PreserveExecutionContext };
+        return new ClientHandler(PathBase, Application)
+        {
+            AllowSynchronousIO = AllowSynchronousIO,
+            PreserveExecutionContext = PreserveExecutionContext
+        };
+    }
+
+    /// <summary>
+    /// Creates a custom <see cref="HttpMessageHandler" /> for processing HTTP requests/responses with custom configuration with the test server.
+    /// </summary>
+    public HttpMessageHandler CreateHandler(Action<HttpContext> additionalContextConfiguration)
+    {
+        return new ClientHandler(PathBase, Application, additionalContextConfiguration)
+        {
+            AllowSynchronousIO = AllowSynchronousIO,
+            PreserveExecutionContext = PreserveExecutionContext
+        };
     }
 
     /// <summary>

src/Hosting/TestHost/test/ClientHandlerTests.cs

@@ -259,6 +259,36 @@ public async Task ServerTrailersSetOnResponseAfterContentRead()
             });
     }
 
+    [Fact]
+    public Task AdditionalConfigurationAllowsSettingConnectionInfo()
+    {
+        var handler = new ClientHandler(PathString.Empty, new InspectingApplication(features =>
+        {
+            Assert.Equal(IPAddress.Parse("1.1.1.1"), features.Get<IHttpConnectionFeature>().RemoteIpAddress);
+        }), context =>
+        {
+            context.Connection.RemoteIpAddress = IPAddress.Parse("1.1.1.1");
+        });
+
+        var httpClient = new HttpClient(handler);
+        return httpClient.GetAsync("https://example.com/A/Path/and/file.txt?and=query");
+    }
+
+    [Fact]
+    public Task AdditionalConfigurationAllowsOverridingDefaultBehavior()
+    {
+        var handler = new ClientHandler(PathString.Empty, new InspectingApplication(features =>
+        {
+            Assert.Equal("?and=something", features.Get<IHttpRequestFeature>().QueryString);
+        }), context =>
+        {
+            context.Request.QueryString = new QueryString("?and=something");
+        });
+
+        var httpClient = new HttpClient(handler);
+        return httpClient.GetAsync("https://example.com/A/Path/and/file.txt?and=query");
+    }
+
     [Fact]
     public async Task ResponseStartAsync()
     {