From 6c4e4db760e5ca52448a579f37b9778c000da036 Mon Sep 17 00:00:00 2001
From: Omer Meshulam <meshulamomer@microsoft.com>
Date: Wed, 24 Dec 2025 22:46:19 +0200
Subject: [PATCH] Added option to pass ClusterModel and DestinationModel to
 IProbingRequestFactory and updated interface, make flow async for
 futureproofing and maximum extensibility

---
 .../Health/ActiveHealthCheckMonitor.cs        |  2 +-
 .../Health/DefaultProbingRequestFactory.cs    |  4 +++
 .../Health/IProbingRequestFactory.cs          | 10 ++++++
 .../Health/ActiveHealthCheckMonitorTests.cs   | 35 ++++++++++++++-----
 4 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/src/ReverseProxy/Health/ActiveHealthCheckMonitor.cs b/src/ReverseProxy/Health/ActiveHealthCheckMonitor.cs
index 6e34e8c36..64ec527e8 100644
--- a/src/ReverseProxy/Health/ActiveHealthCheckMonitor.cs
+++ b/src/ReverseProxy/Health/ActiveHealthCheckMonitor.cs
@@ -175,7 +175,7 @@ private async Task<DestinationProbingResult> ProbeDestinationAsync(ClusterState
         HttpRequestMessage request;
         try
         {
-            request = _probingRequestFactory.CreateRequest(cluster.Model, destination.Model);
+            request = await _probingRequestFactory.CreateRequestAsync(cluster, destination);
         }
         catch (Exception ex)
         {
diff --git a/src/ReverseProxy/Health/DefaultProbingRequestFactory.cs b/src/ReverseProxy/Health/DefaultProbingRequestFactory.cs
index afb62a06a..11ac0b53b 100644
--- a/src/ReverseProxy/Health/DefaultProbingRequestFactory.cs
+++ b/src/ReverseProxy/Health/DefaultProbingRequestFactory.cs
@@ -4,6 +4,7 @@
 using System.Net;
 using System.Net.Http;
 using System.Reflection;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Extensions;
 using Microsoft.Net.Http.Headers;
@@ -39,4 +40,7 @@ public HttpRequestMessage CreateRequest(ClusterModel cluster, DestinationModel d
 
         return request;
     }
+
+    public ValueTask<HttpRequestMessage> CreateRequestAsync(ClusterState cluster, DestinationState destination) =>
+        ValueTask.FromResult(CreateRequest(cluster.Model, destination.Model));
 }
diff --git a/src/ReverseProxy/Health/IProbingRequestFactory.cs b/src/ReverseProxy/Health/IProbingRequestFactory.cs
index 37b334606..80d5302a3 100644
--- a/src/ReverseProxy/Health/IProbingRequestFactory.cs
+++ b/src/ReverseProxy/Health/IProbingRequestFactory.cs
@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Net.Http;
+using System.Threading.Tasks;
 using Yarp.ReverseProxy.Model;
 
 namespace Yarp.ReverseProxy.Health;
@@ -18,4 +19,13 @@ public interface IProbingRequestFactory
     /// <param name="destination">The destination being probed.</param>
     /// <returns>Probing <see cref="HttpRequestMessage"/>.</returns>
     HttpRequestMessage CreateRequest(ClusterModel cluster, DestinationModel destination);
+
+    /// <summary>
+    /// Creates a probing request.
+    /// </summary>
+    /// <param name="cluster">The cluster being probed.</param>
+    /// <param name="destination">The destination being probed.</param>
+    /// <returns>Probing <see cref="HttpRequestMessage"/>.</returns>
+    ValueTask<HttpRequestMessage> CreateRequestAsync(ClusterState cluster, DestinationState destination) =>
+        ValueTask.FromResult(CreateRequest(cluster.Model, destination.Model));
 }
diff --git a/test/ReverseProxy.Tests/Health/ActiveHealthCheckMonitorTests.cs b/test/ReverseProxy.Tests/Health/ActiveHealthCheckMonitorTests.cs
index 889aa8116..a7e4f45f7 100644
--- a/test/ReverseProxy.Tests/Health/ActiveHealthCheckMonitorTests.cs
+++ b/test/ReverseProxy.Tests/Health/ActiveHealthCheckMonitorTests.cs
@@ -58,20 +58,37 @@ public async Task CheckHealthAsync_ActiveHealthCheckIsEnabledForCluster_SendProb
         VerifySentProbeAndResult(cluster2, httpClient2, policy1, new[] { ("https://localhost:20000/cluster2/api/health/", 1), ("https://localhost:20001/cluster2/api/health/", 1) });
     }
 
-    [Fact]
-    public async Task CheckHealthAsync_CustomUserAgentSpecified_UserAgentUnchanged()
+    [Theory]
+    [InlineData(false)] // Test old API (CreateRequest with Models) via default implementation
+    [InlineData(true)]  // Test new API (CreateRequestAsync with State)
+    public async Task CheckHealthAsync_CustomUserAgentSpecified_UserAgentUnchanged(bool overrideAsyncMethod)
     {
         var policy = new Mock<IActiveHealthCheckPolicy>();
         policy.SetupGet(p => p.Name).Returns("policy");
 
         var requestFactory = new Mock<IProbingRequestFactory>();
-        requestFactory.Setup(p => p.CreateRequest(It.IsAny<ClusterModel>(), It.IsAny<DestinationModel>()))
-            .Returns((ClusterModel cluster, DestinationModel destination) =>
-            {
-                var request = new HttpRequestMessage(HttpMethod.Get, "https://localhost:20000/cluster/api/health/");
-                request.Headers.UserAgent.ParseAdd("FooBar/9001");
-                return request;
-            });
+
+        HttpRequestMessage CreateCustomRequest()
+        {
+            var request = new HttpRequestMessage(HttpMethod.Get, "https://localhost:20000/cluster/api/health/");
+            request.Headers.UserAgent.ParseAdd("FooBar/9001");
+            return request;
+        }
+
+        if (overrideAsyncMethod)
+        {
+            requestFactory.Setup(p => p.CreateRequestAsync(It.IsAny<ClusterState>(), It.IsAny<DestinationState>()))
+                .Returns(() => ValueTask.FromResult(CreateCustomRequest()));
+        }
+        else
+        {
+            // Test the old API - the default implementation of CreateRequestAsync should call this
+            requestFactory.Setup(p => p.CreateRequest(It.IsAny<ClusterModel>(), It.IsAny<DestinationModel>()))
+                .Returns(CreateCustomRequest);
+
+            // Use default interface implementation for CreateRequestAsync
+            requestFactory.CallBase = true;
+        }
 
         var options = Options.Create(new ActiveHealthCheckMonitorOptions());
         var monitor = new ActiveHealthCheckMonitor(options, new[] { policy.Object }, requestFactory.Object, new Mock<TimeProvider>().Object, GetLogger());