From 0e72c99a6ad5fdbdbd59ff904405894f64873abb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Hofman?= <michal.hofman@snowflake.com>
Date: Tue, 10 Oct 2023 15:43:55 +0200
Subject: [PATCH] Fixes to flaky tests which fail due to mocked connections
 coming back to the connection pool

---
 Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs | 3 ---
 Snowflake.Data.Tests/Mock/MockRetryUntilRestTimeout.cs | 2 +-
 Snowflake.Data.Tests/Mock/MockSnowflakeDbConnection.cs | 5 +++++
 Snowflake.Data/Client/SnowflakeDbConnection.cs         | 4 ++--
 Snowflake.Data/Core/HttpUtil.cs                        | 5 ++++-
 5 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs b/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs
index 98e9af7c1..8532b3bb0 100755
--- a/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs
+++ b/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs
@@ -229,9 +229,6 @@ public void TestSimpleLargeResultSet()
         }
 
 
-        /*
-         * Disabled to make sure that configuration changes does not cause problems with appveyor
-         */
         [Test, NonParallelizable]
         public void TestUseV1ResultParser()
         {
diff --git a/Snowflake.Data.Tests/Mock/MockRetryUntilRestTimeout.cs b/Snowflake.Data.Tests/Mock/MockRetryUntilRestTimeout.cs
index 8d27a00aa..8397b7145 100644
--- a/Snowflake.Data.Tests/Mock/MockRetryUntilRestTimeout.cs
+++ b/Snowflake.Data.Tests/Mock/MockRetryUntilRestTimeout.cs
@@ -31,7 +31,7 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
                                                               string sid = "")
         {
             // Override the http timeout and set to 1ms to force all http request to timeout and retry
-            message.Properties[BaseRestRequest.HTTP_REQUEST_TIMEOUT_KEY] = TimeSpan.FromMilliseconds(1);
+            message.Properties[BaseRestRequest.HTTP_REQUEST_TIMEOUT_KEY] = TimeSpan.FromTicks(0);;
             return await (base.SendAsync(message, restTimeout, externalCancellationToken).ConfigureAwait(false));
         }
     }
diff --git a/Snowflake.Data.Tests/Mock/MockSnowflakeDbConnection.cs b/Snowflake.Data.Tests/Mock/MockSnowflakeDbConnection.cs
index aa342fe68..0f1758636 100644
--- a/Snowflake.Data.Tests/Mock/MockSnowflakeDbConnection.cs
+++ b/Snowflake.Data.Tests/Mock/MockSnowflakeDbConnection.cs
@@ -92,5 +92,10 @@ private void OnSessionEstablished()
         {
             _connectionState = ConnectionState.Open;
         }
+        
+        protected override bool CanReuseSession(TransactionRollbackStatus transactionRollbackStatus)
+        {
+            return false;
+        }
     }
 }
diff --git a/Snowflake.Data/Client/SnowflakeDbConnection.cs b/Snowflake.Data/Client/SnowflakeDbConnection.cs
index b611589cf..5f66d340c 100755
--- a/Snowflake.Data/Client/SnowflakeDbConnection.cs
+++ b/Snowflake.Data/Client/SnowflakeDbConnection.cs
@@ -37,7 +37,7 @@ public class SnowflakeDbConnection : DbConnection
         // Will fix that in a separated PR though as it's a different issue
         private static Boolean _isArrayBindStageCreated;
 
-        private enum TransactionRollbackStatus
+        protected enum TransactionRollbackStatus
         {
             Undefined, // used to indicate ignored transaction status when pool disabled
             Success,
@@ -232,7 +232,7 @@ public Task CloseAsync(CancellationToken cancellationToken)
             return taskCompletionSource.Task;
         }
 
-        private bool CanReuseSession(TransactionRollbackStatus transactionRollbackStatus)
+        protected virtual bool CanReuseSession(TransactionRollbackStatus transactionRollbackStatus)
         {
             return SnowflakeDbConnectionPool.GetPooling() && 
                    transactionRollbackStatus == TransactionRollbackStatus.Success;
diff --git a/Snowflake.Data/Core/HttpUtil.cs b/Snowflake.Data/Core/HttpUtil.cs
index 0daab0ab6..77d3807ff 100755
--- a/Snowflake.Data/Core/HttpUtil.cs
+++ b/Snowflake.Data/Core/HttpUtil.cs
@@ -364,7 +364,10 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
                         if (!httpTimeout.Equals(Timeout.InfiniteTimeSpan))
                         {
                             childCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-                            childCts.CancelAfter(httpTimeout);
+                            if (httpTimeout.Ticks == 0)
+                                childCts.Cancel();
+                            else
+                                childCts.CancelAfter(httpTimeout);                        
                         }
                         response = await base.SendAsync(requestMessage, childCts == null ?
                             cancellationToken : childCts.Token).ConfigureAwait(false);