Add failover and prelogin count capabilities to test server. Update test case.

Author: mdaigle

src/Microsoft.Data.SqlClient/tests/FunctionalTests/SqlConnectionBasicTests.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Configuration;
 using System.Data;
 using System.Data.Common;
 using System.Reflection;
@@ -16,6 +17,13 @@ namespace Microsoft.Data.SqlClient.Tests
 {
     public class SqlConnectionBasicTests
     {
+        // Reflection
+        public static Assembly systemData = Assembly.GetAssembly(typeof(SqlConnection));
+        public static Type sqlConnection = systemData.GetType("Microsoft.Data.SqlClient.SqlConnection");
+        public static PropertyInfo innerConnectionProperty = sqlConnection.GetProperty("InnerConnection", BindingFlags.NonPublic | BindingFlags.Instance);
+        public static Type sqlInternalConnectionTds = systemData.GetType("Microsoft.Data.SqlClient.SqlInternalConnectionTds");
+        public static PropertyInfo serverProvidedFailoverPartnerProperty = sqlInternalConnectionTds.GetProperty("ServerProvidedFailoverPartner", BindingFlags.NonPublic | BindingFlags.Instance);
+
         [Fact]
         public void ConnectionTest()
         {
@@ -213,18 +221,18 @@ public void ConnectionTestValidCredentialCombination()
         public void TransientFault_IgnoreServerProvidedFailoverPartner_ShouldConnectToUserProvidedPartner()
         {
             // Arrange
-            AppContext.SetSwitch("Microsoft.Data.SqlClient.IgnoreServerProvidedFailoverPartner", true);
+            AppContext.SetSwitch("Switch.Microsoft.Data.SqlClient.IgnoreServerProvidedFailoverPartner", true);
 
             using TestTdsServer failoverServer = TestTdsServer.StartTestServer();
             // Doesn't need to point to a real endpoint, just needs a value specified
-            //TODO: FailoverPartner = "localhost,1234",
+            failoverServer.Arguments.FailoverPartner = "localhost,1234";
 
             var failoverBuilder = new SqlConnectionStringBuilder(failoverServer.ConnectionString);
 
             using TestTdsServer server = TestTdsServer.StartTestServer();
-                    // Set an invalid failover partner to ensure that the connection fails if the
-                    // server provided failover partner is used.
-                    // TODO: FailoverPartner = $"invalidhost",
+            // Set an invalid failover partner to ensure that the connection fails if the
+            // server provided failover partner is used.
+            server.Arguments.FailoverPartner = $"invalidhost";
 
             SqlConnectionStringBuilder builder = new(server.ConnectionString)
             {
@@ -234,18 +242,20 @@ public void TransientFault_IgnoreServerProvidedFailoverPartner_ShouldConnectToUs
                 // Ensure pooling is enabled so that the failover partner information
                 // is persisted in the pool group. If pooling is disabled, the server
                 // provided failover partner will never be used.
-                Pooling = true
+                Pooling = true,
             };
             SqlConnection connection = new(builder.ConnectionString);
 
             // Connect once to the primary to trigger it to send the failover partner
             connection.Open();
-            Assert.Equal("invalidhost", (connection.InnerConnection as SqlInternalConnectionTds)!.ServerProvidedFailoverPartner);
+
+            var innerConnection = innerConnectionProperty.GetValue(connection);
+            var serverProvidedFailoverPartner = serverProvidedFailoverPartnerProperty.GetValue(innerConnection);
+            Assert.Equal("invalidhost", serverProvidedFailoverPartner);
 
             // Close the connection to return it to the pool
             connection.Close();
 
-
             // Act
             // Dispose of the server to trigger a failover
             server.Dispose();

src/Microsoft.Data.SqlClient/tests/tools/TDS/TDS.Servers/GenericTDSServer.cs

@@ -32,16 +32,26 @@ public class GenericTDSServer : ITDSServer
         /// </summary>
         private int _sessionCount = 0;
 
+        /// <summary>
+        /// Counts pre-login requests to the server.
+        /// </summary>
+        private int _preLoginCount = 0;
+
         /// <summary>
         /// Server configuration
         /// </summary>
-        protected TDSServerArguments Arguments { get; set; }
+        public TDSServerArguments Arguments { get; set; }
 
         /// <summary>
         /// Query engine instance
         /// </summary>
         protected QueryEngine Engine { get; set; }
 
+        /// <summary>
+        /// Counts pre-login requests to the server.
+        /// </summary>
+        public int PreLoginCount => _preLoginCount;
+
         /// <summary>
         /// Default constructor
         /// </summary>
@@ -104,6 +114,8 @@ public virtual void CloseSession(ITDSServerSession session)
         /// </summary>
         public virtual TDSMessageCollection OnPreLoginRequest(ITDSServerSession session, TDSMessage request)
         {
+            Interlocked.Increment(ref _preLoginCount);
+
             // Inflate pre-login request from the message
             TDSPreLoginToken preLoginRequest = request[0] as TDSPreLoginToken;
 
@@ -544,6 +556,16 @@ protected virtual TDSMessageCollection OnAuthenticationCompleted(ITDSServerSessi
                 responseMessage.Add(featureExtActToken);
             }
 
+            if (!string.IsNullOrEmpty(Arguments.FailoverPartner))
+            {
+                envChange = new TDSEnvChangeToken(TDSEnvChangeTokenType.RealTimeLogShipping, Arguments.FailoverPartner);
+
+                // Log response
+                TDSUtilities.Log(Arguments.Log, "Response", envChange);
+
+                responseMessage.Add(envChange);
+            }
+
             // Create DONE token
             TDSDoneToken doneToken = new TDSDoneToken(TDSDoneTokenStatusType.Final);
 

src/Microsoft.Data.SqlClient/tests/tools/TDS/TDS.Servers/TDSServerArguments.cs

@@ -69,6 +69,11 @@ public class TDSServerArguments
         /// </summary>
         public X509Certificate EncryptionCertificate { get; set; }
 
+        /// <summary>
+        /// Specifies the failover partner server name and port
+        /// </summary>
+        public string FailoverPartner { get; set; } = string.Empty;
+
         /// <summary>
         /// Initialization constructor
         /// </summary>