Add tests for timeout evaluation using switch.

Author: mdaigle

src/Microsoft.Data.SqlClient/src/Microsoft/Data/ProviderBase/TimeoutTimer.cs

@@ -206,7 +206,7 @@ internal int MillisecondsRemainingInt
         /// timer was created. Used by <see cref="Reset"/> to restore the original
         /// expiration window.
         /// </summary>
-        private long OriginalTicks { get; }
+        internal long OriginalTicks { get; }
 
         /// <summary>
         /// Gets the <see cref="TimeProvider"/> used by this timer. Exposed for

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/Connection/SqlConnectionInternal.cs

@@ -403,14 +403,7 @@ internal SqlConnectionInternal(
                 // If we want to consider pool operations against the overall connect timeout, 
                 // use the provided timeout. Otherwise, start a fresh timeout to receive the full
                 // connect timeout.
-                if (LocalAppContextSwitches.UseOverallConnectTimeoutForPoolWait)
-                {
-                    _timeout = timeout;
-                }
-                else
-                {
-                    _timeout = TimeoutTimer.StartNew(TimeSpan.FromSeconds(connectionOptions.ConnectTimeout));
-                }
+                _timeout = ResolveLoginTimeout(timeout, connectionOptions.ConnectTimeout);
 
                 // If transient fault handling is enabled then we can retry the login up to the
                 // ConnectRetryCount.
@@ -777,6 +770,25 @@ private SqlInternalTransaction AvailableInternalTransaction
         // @TODO: Make private field.
         private bool IsAzureSqlConnection { get; set; }
 
+        internal TimeoutTimer Timeout => _timeout;
+
+        /// <summary>
+        /// Selects the <see cref="TimeoutTimer"/> that governs the login phase based on
+        /// <see cref="LocalAppContextSwitches.UseOverallConnectTimeoutForPoolWait"/>.
+        /// </summary>
+        /// <remarks>
+        /// When the switch is enabled the caller-supplied <paramref name="callerTimeout"/>
+        /// is returned as-is so any time already consumed (e.g., waiting for the pool) counts
+        /// against the overall ConnectTimeout. When disabled, a fresh timer is started from
+        /// <paramref name="connectTimeoutSeconds"/>, preserving legacy behavior. Extracted so
+        /// this branch can be unit-tested without standing up a real connection.
+        /// </remarks>
+        internal static TimeoutTimer ResolveLoginTimeout(TimeoutTimer callerTimeout, int connectTimeoutSeconds)
+        {
+            return LocalAppContextSwitches.UseOverallConnectTimeoutForPoolWait
+                ? callerTimeout
+                : TimeoutTimer.StartNew(TimeSpan.FromSeconds(connectTimeoutSeconds));
+        }
         #endregion
 
         #region Public and Internal Methods

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/WaitHandleDbConnectionPool.cs

@@ -875,38 +875,41 @@ private void WaitForPendingOpen()
             } while (_pendingOpens.TryPeek(out next));
         }
 
+        /// <summary>
+        /// Resolves the <c>WaitHandle.WaitAny</c> timeout (milliseconds) for a synchronous
+        /// pool acquire.
+        /// </summary>
+        /// <remarks>
+        /// When <see cref="LocalAppContextSwitches.UseOverallConnectTimeoutForPoolWait"/>
+        /// is enabled the caller's remaining <see cref="TimeoutTimer"/> budget is used so
+        /// the pool wait participates in the overall ConnectTimeout. Otherwise the legacy
+        /// behavior is preserved: the static pool <c>CreationTimeout</c> is used, with
+        /// <c>0</c> mapped to <see cref="Timeout.Infinite"/>. Extracted so this branch can
+        /// be unit-tested without timing-based assertions.
+        /// </remarks>
+        internal static uint ResolvePoolWaitTimeoutMs(TimeoutTimer timeout, int creationTimeoutMs)
+        {
+            if (LocalAppContextSwitches.UseOverallConnectTimeoutForPoolWait)
+            {
+                return timeout.IsInfinite
+                    ? unchecked((uint)Timeout.Infinite)
+                    : (uint)timeout.MillisecondsRemainingInt;
+            }
+
+            uint legacy = (uint)creationTimeoutMs;
+            return legacy == 0 ? unchecked((uint)Timeout.Infinite) : legacy;
+        }
+
         public bool TryGetConnection(DbConnection owningObject, TaskCompletionSource<DbConnectionInternal> taskCompletionSource, TimeoutTimer timeout, out DbConnectionInternal connection)
         {
             uint waitForMultipleObjectsTimeout = 0;
             bool allowCreate = false;
 
             if (taskCompletionSource == null)
             {
-                if (LocalAppContextSwitches.UseOverallConnectTimeoutForPoolWait)
-                {
-                    // Use the caller's remaining timeout budget (rather than the static
-                    // CreationTimeout) so synchronous pool waits respect any time already
-                    // consumed earlier in the open path and don't exceed the overall
-                    // ConnectTimeout.
-                    waitForMultipleObjectsTimeout = timeout.IsInfinite
-                        ? unchecked((uint)Timeout.Infinite)
-                        : (uint)timeout.MillisecondsRemainingInt;
-                }
-                else
-                {
-                    waitForMultipleObjectsTimeout = (uint)CreationTimeout;
-
-                    // Set the wait timeout to INFINITE (-1) if the pool CreationTimeout is 0.
-                    if (waitForMultipleObjectsTimeout == 0)
-                    {
-                        waitForMultipleObjectsTimeout = unchecked((uint)Timeout.Infinite);
-                    }
-                }
-
+                waitForMultipleObjectsTimeout = ResolvePoolWaitTimeoutMs(timeout, CreationTimeout);
                 allowCreate = true;
-            }
-
-            if (State is not Running)
+            }            if (State is not Running)
             {
                 SqlClientEventSource.Log.TryPoolerTraceEvent("<prov.DbConnectionPool.GetConnection|RES|CPOOL> {0}, DbConnectionInternal State != Running.", Id);
                 connection = null;

src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/WaitHandleDbConnectionPoolBudgetTest.cs

@@ -4,10 +4,12 @@
 
 using System;
 using System.Data.Common;
+using System.Threading;
 using System.Threading.Tasks;
 using System.Transactions;
 using Microsoft.Data.ProviderBase;
 using Microsoft.Data.SqlClient.ConnectionPool;
+using Microsoft.Data.SqlClient.Tests.Common;
 using Microsoft.Extensions.Time.Testing;
 using Xunit;
 
@@ -65,73 +67,110 @@ private WaitHandleDbConnectionPool CreatePool(
 
     /// <summary>
     /// Verifies that the <see cref="TimeoutTimer"/> the pool hands to the
-    /// connection factory on the synchronous path reports a reduced
-    /// remaining-time budget when the timer's clock has advanced before the
-    /// pool was entered. Mirrors
+    /// connection factory reports a reduced remaining-time budget when the
+    /// timer's clock has advanced before the pool was entered. Both the
+    /// synchronous (<c>taskCompletionSource == null</c>) and asynchronous
+    /// paths must forward the caller's already-advanced timer rather than
+    /// constructing a fresh one from <c>CreationTimeout</c>. Mirrors
     /// <c>ChannelDbConnectionPoolTest.GetConnection_TimeoutTimerReflectsPoolWaitTime</c>.
     /// </summary>
-    [Fact]
-    public void GetConnection_Sync_TimeoutTimerReflectsTimeAlreadyConsumed()
+    [Theory]
+    [InlineData(false)] // sync
+    [InlineData(true)]  // async
+    public async Task GetConnection_TimeoutTimerReflectsTimeAlreadyConsumed(bool async)
     {
-        // Arrange: a capturing factory and a fake-time-backed timer with a
+        // Arrange: capturing factory and a fake-time-backed timer with a
         // 30-second budget anchored at virtual time t = 0.
         var factory = new MockSqlConnectionFactory();
         var pool = CreatePool(factory);
         var owner = new SqlConnection("Timeout=30");
         var fakeTime = new FakeTimeProvider();
         TimeoutTimer timer = TimeoutTimer.StartNew(TimeSpan.FromSeconds(30), fakeTime);
+        TaskCompletionSource<DbConnectionInternal>? tcs =
+            async ? new TaskCompletionSource<DbConnectionInternal>() : null;
 
         // Act: simulate 5s of budget consumed elsewhere (e.g., higher-level
-        // Open() work) before the pool is entered.
+        // Open() work) before the pool is entered, then request a connection.
         fakeTime.Advance(TimeSpan.FromSeconds(5));
         bool completed = pool.TryGetConnection(
             owner,
-            taskCompletionSource: null,
+            tcs,
             timer,
             out DbConnectionInternal? connection);
 
+        if (async)
+        {
+            // Bound the await so a regression in the pool can't hang the suite.
+            Task winner = await Task.WhenAny(tcs!.Task, Task.Delay(TimeSpan.FromSeconds(30)));
+            Assert.Same(tcs.Task, winner);
+            connection = await tcs.Task;
+        }
+        else
+        {
+            Assert.True(completed);
+        }
+
         // Assert: factory received the same timer, and it reports the reduced
         // 25-second remaining budget rather than the original 30s or the
         // pool's static 15s CreationTimeout.
-        Assert.True(completed);
         Assert.NotNull(connection);
         Assert.Same(timer, factory.CapturedTimeout);
         Assert.Equal(25_000, factory.CapturedTimeout!.MillisecondsRemainingInt);
     }
 
     /// <summary>
-    /// Async counterpart of <see cref="GetConnection_Sync_TimeoutTimerReflectsTimeAlreadyConsumed"/>.
-    /// Verifies that the async pool path also forwards the caller's
-    /// already-advanced <see cref="TimeoutTimer"/> to the factory.
+    /// Identifies which kind of caller-supplied <see cref="TimeoutTimer"/> a
+    /// parameterized test should construct. Used because
+    /// <see cref="InlineDataAttribute"/> cannot carry a live
+    /// <see cref="TimeoutTimer"/> instance.
+    /// </summary>
+    public enum TimerKind
+    {
+        Expired,
+        Infinite,
+    }
+
+    /// <summary>
+    /// Verifies the resolution matrix for the synchronous
+    /// <c>WaitHandle.WaitAny</c> timeout:
+    /// <list type="bullet">
+    ///   <item>switch ON  → use the caller timer's remaining budget
+    ///   (expired → 0, infinite → <see cref="Timeout.Infinite"/>);</item>
+    ///   <item>switch OFF → ignore the caller timer and use
+    ///   <c>CreationTimeout</c>, treating <c>0</c> as
+    ///   <see cref="Timeout.Infinite"/> per legacy behavior.</item>
+    /// </list>
     /// </summary>
-    [Fact]
-    public async Task GetConnection_Async_TimeoutTimerReflectsTimeAlreadyConsumed()
+    [Theory]
+    [InlineData(true,  TimerKind.Expired,  5_000, 0u)]
+    [InlineData(true,  TimerKind.Infinite, 5_000, unchecked((uint)Timeout.Infinite))]
+    [InlineData(false, TimerKind.Expired,  1_500, 1_500u)]
+    [InlineData(false, TimerKind.Expired,  0,     unchecked((uint)Timeout.Infinite))]
+    public void ResolvePoolWaitTimeoutMs_ReturnsExpected(
+        bool switchEnabled,
+        TimerKind timerKind,
+        int creationTimeoutMs,
+        uint expected)
     {
         // Arrange
-        var factory = new MockSqlConnectionFactory();
-        var pool = CreatePool(factory);
-        var owner = new SqlConnection("Timeout=30");
-        var fakeTime = new FakeTimeProvider();
-        TimeoutTimer timer = TimeoutTimer.StartNew(TimeSpan.FromSeconds(30), fakeTime);
-        var tcs = new TaskCompletionSource<DbConnectionInternal>();
+        using LocalAppContextSwitchesHelper switches = new()
+        {
+            UseOverallConnectTimeoutForPoolWait = switchEnabled,
+        };
+        TimeoutTimer timer = timerKind switch
+        {
+            TimerKind.Expired => TimeoutTimer.StartExpired(),
+            TimerKind.Infinite => TimeoutTimer.StartNew(TimeSpan.Zero),
+            _ => throw new ArgumentOutOfRangeException(nameof(timerKind)),
+        };
 
-        // Act: 5s consumed before entering the pool, then an async request.
-        fakeTime.Advance(TimeSpan.FromSeconds(5));
-        pool.TryGetConnection(
-            owner,
-            taskCompletionSource: tcs,
+        // Act
+        uint result = WaitHandleDbConnectionPool.ResolvePoolWaitTimeoutMs(
             timer,
-            out DbConnectionInternal? connection);
-
-        // Bound the await so a regression in the pool can't hang the suite.
-        Task completed = await Task.WhenAny(tcs.Task, Task.Delay(TimeSpan.FromSeconds(30)));
-        Assert.Same(tcs.Task, completed);
-        DbConnectionInternal result = await tcs.Task;
+            creationTimeoutMs);
 
-        // Assert: factory got the caller's timer with the reduced budget.
-        Assert.NotNull(result);
-        Assert.Same(timer, factory.CapturedTimeout);
-        Assert.Equal(25_000, factory.CapturedTimeout!.MillisecondsRemainingInt);
+        // Assert
+        Assert.Equal(expected, result);
     }
 
     /// <summary>

src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/SqlConnectionInternalTimeoutTests.cs

@@ -0,0 +1,61 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using Microsoft.Data.ProviderBase;
+using Microsoft.Data.SqlClient.Connection;
+using Microsoft.Data.SqlClient.Tests.Common;
+using Xunit;
+
+namespace Microsoft.Data.SqlClient.UnitTests;
+
+/// <summary>
+/// Verifies how <see cref="SqlConnectionInternal"/> selects the
+/// <see cref="TimeoutTimer"/> that governs the login phase, driven by the
+/// <c>UseOverallConnectTimeoutForPoolWait</c> AppContext switch.
+///
+/// When the switch is enabled the connection must reuse the caller-supplied
+/// timer as-is so the remaining budget (already reflecting any time spent
+/// waiting for the pool) is honored during login. When disabled it must
+/// construct a fresh timer from <c>ConnectTimeout</c>, preserving legacy
+/// behavior. Asserting against the extracted
+/// <see cref="SqlConnectionInternal.ResolveLoginTimeout"/> helper keeps this
+/// branch coverage free of any real network connection.
+/// </summary>
+public class SqlConnectionInternalTimeoutTests
+{
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public void ResolveLoginTimeout_HonorsSwitch(bool switchEnabled)
+    {
+        // Arrange
+        using LocalAppContextSwitchesHelper switches = new()
+        {
+            UseOverallConnectTimeoutForPoolWait = switchEnabled,
+        };
+        const int ConnectTimeoutSeconds = 30;
+        TimeoutTimer callerTimeout = TimeoutTimer.StartNew(TimeSpan.FromSeconds(ConnectTimeoutSeconds));
+
+        // Act
+        TimeoutTimer resolved = SqlConnectionInternal.ResolveLoginTimeout(
+            callerTimeout,
+            ConnectTimeoutSeconds);
+
+        // Assert
+        if (switchEnabled)
+        {
+            // Switch on: caller's timer must flow through unchanged so any
+            // time already consumed counts against the overall budget.
+            Assert.Same(callerTimeout, resolved);
+        }
+        else
+        {
+            // Switch off (legacy): a fresh timer must be started from
+            // ConnectTimeout, independent of the caller's timer.
+            Assert.NotSame(callerTimeout, resolved);
+            Assert.Equal(TimeSpan.FromSeconds(ConnectTimeoutSeconds).Ticks, resolved.OriginalTicks);
+        }
+    }
+}