From d65d22b1ec6cbce11cd3db56c710cacaeba4d7d4 Mon Sep 17 00:00:00 2001
From: Artak Mkrtchyan <artakm@microsoft.com>
Date: Wed, 26 Feb 2025 13:47:14 -0800
Subject: [PATCH] Add a new UseKestel API to the WebApplicationFactory

---
 .../Microsoft.AspNetCore.Mvc.Testing.csproj   |   1 +
 .../Mvc.Testing/src/PublicAPI.Unshipped.txt   |   4 +
 src/Mvc/Mvc.Testing/src/Resources.resx        |   5 +-
 .../Mvc.Testing/src/WebApplicationFactory.cs  | 109 ++++++++++++++++--
 4 files changed, 106 insertions(+), 13 deletions(-)

diff --git a/src/Mvc/Mvc.Testing/src/Microsoft.AspNetCore.Mvc.Testing.csproj b/src/Mvc/Mvc.Testing/src/Microsoft.AspNetCore.Mvc.Testing.csproj
index 3201fd08940b..76118e0bb154 100644
--- a/src/Mvc/Mvc.Testing/src/Microsoft.AspNetCore.Mvc.Testing.csproj
+++ b/src/Mvc/Mvc.Testing/src/Microsoft.AspNetCore.Mvc.Testing.csproj
@@ -13,6 +13,7 @@
   <ItemGroup>
     <Reference Include="Microsoft.AspNetCore.TestHost" />
     <Reference Include="Microsoft.AspNetCore.Mvc.Core" />
+    <Reference Include="Microsoft.AspNetCore.Server.Kestrel" />
     <Reference Include="Microsoft.Extensions.DependencyModel" />
     <Reference Include="Microsoft.Extensions.Hosting" />
     <Reference Include="Microsoft.Extensions.HostFactoryResolver.Sources" />
diff --git a/src/Mvc/Mvc.Testing/src/PublicAPI.Unshipped.txt b/src/Mvc/Mvc.Testing/src/PublicAPI.Unshipped.txt
index 7dc5c58110bf..d8f82d0db79f 100644
--- a/src/Mvc/Mvc.Testing/src/PublicAPI.Unshipped.txt
+++ b/src/Mvc/Mvc.Testing/src/PublicAPI.Unshipped.txt
@@ -1 +1,5 @@
 #nullable enable
+*REMOVED*Microsoft.AspNetCore.Mvc.Testing.WebApplicationFactory<TEntryPoint>.Server.get -> Microsoft.AspNetCore.TestHost.TestServer!
+Microsoft.AspNetCore.Mvc.Testing.WebApplicationFactory<TEntryPoint>.Server.get -> Microsoft.AspNetCore.TestHost.TestServer?
+Microsoft.AspNetCore.Mvc.Testing.WebApplicationFactory<TEntryPoint>.UseKestrel() -> void
+virtual Microsoft.AspNetCore.Mvc.Testing.WebApplicationFactory<TEntryPoint>.CreateHandler() -> System.Net.Http.HttpMessageHandler!
diff --git a/src/Mvc/Mvc.Testing/src/Resources.resx b/src/Mvc/Mvc.Testing/src/Resources.resx
index 84ce491ca5ed..abfb9fcfe3b3 100644
--- a/src/Mvc/Mvc.Testing/src/Resources.resx
+++ b/src/Mvc/Mvc.Testing/src/Resources.resx
@@ -1,4 +1,4 @@
-<?xml version="1.0" encoding="utf-8"?>
+<?xml version="1.0" encoding="utf-8"?>
 <root>
   <!-- 
     Microsoft ResX Schema 
@@ -126,4 +126,7 @@
   <data name="MissingDepsFile" xml:space="preserve">
     <value>Can't find '{0}'. This file is required for functional tests to run properly. There should be a copy of the file on your source project bin folder. If that is not the case, make sure that the property PreserveCompilationContext is set to true on your project file. E.g '&lt;PreserveCompilationContext&gt;true&lt;/PreserveCompilationContext&gt;'. For functional tests to work they need to either run from the build output folder or the {1} file from your application's output directory must be copied to the folder where the tests are running on. A common cause for this error is having shadow copying enabled when the tests run.</value>
   </data>
+  <data name="ServerNotInitialized" xml:space="preserve">
+    <value>Server hasn't been initialized yet. Plase intialized the server first before trying to create a client.</value>
+  </data>
 </root>
\ No newline at end of file
diff --git a/src/Mvc/Mvc.Testing/src/WebApplicationFactory.cs b/src/Mvc/Mvc.Testing/src/WebApplicationFactory.cs
index dd43bdac2e7c..f48be3b5b94b 100644
--- a/src/Mvc/Mvc.Testing/src/WebApplicationFactory.cs
+++ b/src/Mvc/Mvc.Testing/src/WebApplicationFactory.cs
@@ -1,7 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Net.Http;
 using System.Reflection;
@@ -9,6 +8,7 @@
 using System.Text.Json.Serialization;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.Hosting.Server.Features;
 using Microsoft.AspNetCore.TestHost;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
@@ -24,11 +24,13 @@ namespace Microsoft.AspNetCore.Mvc.Testing;
 /// Typically the Startup or Program classes can be used.</typeparam>
 public partial class WebApplicationFactory<TEntryPoint> : IDisposable, IAsyncDisposable where TEntryPoint : class
 {
+    private bool _useKestrel;
     private bool _disposed;
     private bool _disposedAsync;
     private TestServer? _server;
     private IHost? _host;
     private Action<IWebHostBuilder> _configuration;
+    private IWebHost? _webHost;
     private readonly List<HttpClient> _clients = new();
     private readonly List<WebApplicationFactory<TEntryPoint>> _derivedFactories = new();
 
@@ -71,7 +73,7 @@ public WebApplicationFactory()
     /// <summary>
     /// Gets the <see cref="TestServer"/> created by this <see cref="WebApplicationFactory{TEntryPoint}"/>.
     /// </summary>
-    public TestServer Server
+    public TestServer? Server
     {
         get
         {
@@ -88,7 +90,12 @@ public virtual IServiceProvider Services
         get
         {
             EnsureServer();
-            return _host?.Services ?? _server.Host.Services;
+            if (_useKestrel)
+            {
+                return _webHost!.Services;
+            }
+
+            return _host?.Services ?? _server!.Host.Services;
         }
     }
 
@@ -136,10 +143,24 @@ internal virtual WebApplicationFactory<TEntryPoint> WithWebHostBuilderCore(Actio
         return factory;
     }
 
-    [MemberNotNull(nameof(_server))]
+    /// <summary>
+    /// Configures the factory to use Kestrel as the server.
+    /// </summary>
+    public void UseKestrel()
+    {
+        _useKestrel = true;
+    }
+
+    private static IWebHost CreateKestrelServer(IWebHostBuilder builder)
+    {
+        var host = builder.UseKestrel().Build();
+        host.Start();
+        return host;
+    }
+
     private void EnsureServer()
     {
-        if (_server != null)
+        if (_server != null || _webHost != null)
         {
             return;
         }
@@ -197,21 +218,37 @@ private void EnsureServer()
         {
             SetContentRoot(builder);
             _configuration(builder);
-            _server = CreateServer(builder);
+            if (_useKestrel)
+            {
+                _webHost = CreateKestrelServer(builder);
+            }
+            else
+            {
+                _server = CreateServer(builder);
+            }
         }
     }
 
-    [MemberNotNull(nameof(_server))]
     private void ConfigureHostBuilder(IHostBuilder hostBuilder)
     {
         hostBuilder.ConfigureWebHost(webHostBuilder =>
         {
             SetContentRoot(webHostBuilder);
             _configuration(webHostBuilder);
-            webHostBuilder.UseTestServer();
+            if (!_useKestrel)
+            {
+                webHostBuilder.UseTestServer();
+            }
+            else
+            {
+                webHostBuilder.UseKestrel();
+            }
         });
         _host = CreateHost(hostBuilder);
-        _server = (TestServer)_host.Services.GetRequiredService<IServer>();
+        if (!_useKestrel)
+        {
+            _server = (TestServer)_host.Services.GetRequiredService<IServer>();
+        }
     }
 
     private void SetContentRoot(IWebHostBuilder builder)
@@ -481,7 +518,19 @@ public HttpClient CreateDefaultClient(params DelegatingHandler[] handlers)
         HttpClient client;
         if (handlers == null || handlers.Length == 0)
         {
-            client = _server.CreateClient();
+            if (_useKestrel)
+            {
+                client = new HttpClient();
+            }
+            else
+            {
+                if (_server is null)
+                {
+                    throw new InvalidOperationException(Resources.ServerNotInitialized);
+                }
+
+                client = _server.CreateClient();
+            }
         }
         else
         {
@@ -490,7 +539,7 @@ public HttpClient CreateDefaultClient(params DelegatingHandler[] handlers)
                 handlers[i - 1].InnerHandler = handlers[i];
             }
 
-            var serverHandler = _server.CreateHandler();
+            var serverHandler = CreateHandler();
             handlers[^1].InnerHandler = serverHandler;
 
             client = new HttpClient(handlers[0]);
@@ -503,6 +552,25 @@ public HttpClient CreateDefaultClient(params DelegatingHandler[] handlers)
         return client;
     }
 
+    /// <summary>
+    /// Creates a custom <see cref="HttpMessageHandler" /> for processing HTTP requests/responses with the test server.
+    /// </summary>
+
+    protected virtual HttpMessageHandler CreateHandler()
+    {
+        if (_useKestrel)
+        {
+            return new HttpClientHandler();
+        }
+
+        if (_server is null)
+        {
+            throw new InvalidOperationException(Resources.ServerNotInitialized);
+        }
+
+        return _server.CreateHandler();
+    }
+
     /// <summary>
     /// Configures <see cref="HttpClient"/> instances created by this <see cref="WebApplicationFactory{TEntryPoint}"/>.
     /// </summary>
@@ -511,7 +579,24 @@ protected virtual void ConfigureClient(HttpClient client)
     {
         ArgumentNullException.ThrowIfNull(client);
 
-        client.BaseAddress = new Uri("http://localhost");
+        if (_useKestrel)
+        {
+            if (_webHost is null)
+            {
+                throw new InvalidOperationException(Resources.ServerNotInitialized);
+            }
+
+            // The server should be started at this point already.
+            var serverAddressFeature = _webHost.ServerFeatures.Get<IServerAddressesFeature>();
+            if (serverAddressFeature?.Addresses.Count > 0)
+            {
+                client.BaseAddress = new Uri(serverAddressFeature.Addresses.Last());
+            }
+        }
+        else
+        {
+            client.BaseAddress = new Uri("http://localhost");
+        }
     }
 
     /// <summary>