From 0fd4856de7c07ca5f8b368e2b88a21cc425cfa53 Mon Sep 17 00:00:00 2001
From: sfc-gh-ext-simba-lf <lfarol@magnitude.com>
Date: Fri, 9 Feb 2024 10:08:56 -0800
Subject: [PATCH] SNOW-916949: Add tests

---
 .../IntegrationTests/SFConnectionIT.cs        | 71 +++++++++++++++++++
 .../Mock/MockOktaRetryMaxTimeout.cs           | 55 ++++++++++++++
 .../UnitTests/HttpUtilTest.cs                 | 18 +++++
 Snowflake.Data.Tests/UnitTests/SFOktaTest.cs  | 12 +++-
 4 files changed, 154 insertions(+), 2 deletions(-)
 create mode 100644 Snowflake.Data.Tests/Mock/MockOktaRetryMaxTimeout.cs

diff --git a/Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs b/Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs
index 5a4976162..2470068eb 100644
--- a/Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs
+++ b/Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs
@@ -817,6 +817,40 @@ public void TestOktaConnection()
             }
         }
 
+        [Test]
+        public void TestOktaConnectionUntilMaxTimeout()
+        {
+            var expectedMaxRetryCount = 15;
+            var expectedMaxConnectionTimeout = 450;
+            var mockRestRequester = new MockOktaRetryMaxTimeout(expectedMaxRetryCount, expectedMaxConnectionTimeout);
+            using (DbConnection conn = new MockSnowflakeDbConnection(mockRestRequester))
+            {
+                try
+                {
+                    conn.ConnectionString
+                        = ConnectionStringWithoutAuth
+                        + String.Format(
+                            ";authenticator={0};user={1};password={2};MAXHTTPRETRIES={3};RETRY_TIMEOUT={4};",
+                            testConfig.oktaUrl,
+                            testConfig.oktaUser,
+                            testConfig.oktaPassword,
+                            expectedMaxRetryCount,
+                            expectedMaxConnectionTimeout);
+                    conn.Open();
+                    Assert.Fail();
+                }
+                catch (Exception e)
+                {
+                    Assert.IsInstanceOf<SnowflakeDbException>(e);
+                    Assert.AreEqual(SFError.INTERNAL_ERROR.GetAttribute<SFErrorAttr>().errorCode, ((SnowflakeDbException)e).ErrorCode);
+                    Assert.IsTrue(e.Message.Contains(
+                        $"The retry count has reached its limit of {expectedMaxRetryCount} and " +
+                        $"the timeout elapsed has reached its limit of {expectedMaxConnectionTimeout} " +
+                        "while trying to authenticate through Okta"));
+                }
+            }
+        }
+
         [Test]
         [Ignore("This test requires manual setup and therefore cannot be run in CI")]
         public void TestOkta2ConnectionsFollowingEachOther()
@@ -2056,6 +2090,43 @@ public void TestExplicitTransactionOperationsTracked()
                 Assert.AreEqual(false, conn.HasActiveExplicitTransaction());
             }
         }
+
+
+        [Test]
+        public void TestAsyncOktaConnectionUntilMaxTimeout()
+        {
+            var expectedMaxRetryCount = 15;
+            var expectedMaxConnectionTimeout = 450;
+            var mockRestRequester = new MockOktaRetryMaxTimeout(expectedMaxRetryCount, expectedMaxConnectionTimeout);
+            using (DbConnection conn = new MockSnowflakeDbConnection(mockRestRequester))
+            {
+                Task connectTask = null;
+                try
+                {
+                    conn.ConnectionString
+                        = ConnectionStringWithoutAuth
+                        + String.Format(
+                            ";authenticator={0};user={1};password={2};MAXHTTPRETRIES={3};RETRY_TIMEOUT={4};",
+                            testConfig.oktaUrl,
+                            testConfig.oktaUser,
+                            testConfig.oktaPassword,
+                            expectedMaxRetryCount,
+                            expectedMaxConnectionTimeout);
+                    connectTask = conn.OpenAsync(CancellationToken.None);
+                    connectTask.Wait();
+                    Assert.Fail();
+                }
+                catch (Exception e)
+                {
+                    Assert.IsInstanceOf<SnowflakeDbException>(e.InnerException);
+                    Assert.AreEqual(SFError.INTERNAL_ERROR.GetAttribute<SFErrorAttr>().errorCode, ((SnowflakeDbException)e.InnerException).ErrorCode);
+                    Assert.IsTrue(e.InnerException.InnerException.Message.Contains(
+                        $"The retry count has reached its limit of {expectedMaxRetryCount} and " +
+                        $"the timeout elapsed has reached its limit of {expectedMaxConnectionTimeout} " +
+                        "while trying to authenticate through Okta"));
+                }
+            }
+        }
     }
 }
 
diff --git a/Snowflake.Data.Tests/Mock/MockOktaRetryMaxTimeout.cs b/Snowflake.Data.Tests/Mock/MockOktaRetryMaxTimeout.cs
new file mode 100644
index 000000000..20f7c573d
--- /dev/null
+++ b/Snowflake.Data.Tests/Mock/MockOktaRetryMaxTimeout.cs
@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
+ */
+
+using Newtonsoft.Json;
+using Snowflake.Data.Core;
+using Snowflake.Data.Core.Authenticator;
+using System;
+using System.Net.Http;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Snowflake.Data.Tests.Mock
+{
+
+    class MockOktaRetryMaxTimeout : RestRequester, IMockRestRequester
+    {
+        internal bool _forceTimeoutForNonLoginRequestsOnly = false;
+        internal int _maxRetryCount;
+        internal int _maxRetryTimeout;
+
+        public MockOktaRetryMaxTimeout(int maxRetryCount, int maxRetryTimeout) : base(null)
+        {
+            _maxRetryCount = maxRetryCount;
+            _maxRetryTimeout = maxRetryTimeout;
+        }
+
+        public void setHttpClient(HttpClient httpClient)
+        {
+            _HttpClient = httpClient;
+        }
+
+        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage message,
+                                                              TimeSpan restTimeout,
+                                                              CancellationToken externalCancellationToken,
+                                                              string sid = "")
+        {
+            if (HttpUtil.IsOktaSSORequest(message.RequestUri.Host, message.RequestUri.AbsolutePath))
+            {
+                var mockContent = new StringContent(JsonConvert.SerializeObject("<form=error}"), Encoding.UTF8, "application/json");
+                mockContent.Headers.Add(OktaAuthenticator.RetryCountHeader, _maxRetryCount.ToString());
+                mockContent.Headers.Add(OktaAuthenticator.TimeoutElapsedHeader, _maxRetryTimeout.ToString());
+
+                return new HttpResponseMessage
+                {
+                    StatusCode = System.Net.HttpStatusCode.OK,
+                    Content = mockContent
+                };
+            }
+
+            return await base.SendAsync(message, restTimeout, externalCancellationToken).ConfigureAwait(false);
+        }
+    }
+}
diff --git a/Snowflake.Data.Tests/UnitTests/HttpUtilTest.cs b/Snowflake.Data.Tests/UnitTests/HttpUtilTest.cs
index bf8063a7f..668db5f91 100644
--- a/Snowflake.Data.Tests/UnitTests/HttpUtilTest.cs
+++ b/Snowflake.Data.Tests/UnitTests/HttpUtilTest.cs
@@ -58,6 +58,24 @@ public void TestIsLoginUrl(string requestUrl, bool expectedIsLoginEndpoint)
             Assert.AreEqual(expectedIsLoginEndpoint, isLoginEndpoint);
         }
 
+        // Parameters: request url, expected value
+        [TestCase("https://dev.okta.com/sso/saml", true)]
+        [TestCase("https://test.snowflakecomputing.com/session/v1/login-request", false)]
+        [TestCase("https://test.snowflakecomputing.com/session/authenticator-request", false)]
+        [TestCase("https://test.snowflakecomputing.com/session/token-request", false)]
+        [Test]
+        public void TestIsOktaSSORequest(string requestUrl, bool expectedIsOktaSSORequest)
+        {
+            // given
+            var uri = new Uri(requestUrl);
+
+            // when
+            bool isOktaSSORequest = HttpUtil.IsOktaSSORequest(uri.Host, uri.AbsolutePath);
+
+            // then
+            Assert.AreEqual(expectedIsOktaSSORequest, isOktaSSORequest);
+        }
+
         // Parameters: time in seconds
         [TestCase(4)]
         [TestCase(8)]
diff --git a/Snowflake.Data.Tests/UnitTests/SFOktaTest.cs b/Snowflake.Data.Tests/UnitTests/SFOktaTest.cs
index 5de0c6c06..8b5c9dc55 100644
--- a/Snowflake.Data.Tests/UnitTests/SFOktaTest.cs
+++ b/Snowflake.Data.Tests/UnitTests/SFOktaTest.cs
@@ -1,6 +1,7 @@
 using NUnit.Framework;
 using Snowflake.Data.Client;
 using Snowflake.Data.Core;
+using Snowflake.Data.Core.Authenticator;
 using System.Net.Http;
 
 namespace Snowflake.Data.Tests.UnitTests
@@ -33,6 +34,10 @@ public void TestSsoTokenUrlMismatch()
         [Test]
         public void TestMissingPostbackUrl()
         {
+            const int retryCount = 1;
+            const int retryTimeout = 1;
+            noPostbackContent.Headers.Add(OktaAuthenticator.RetryCountHeader, retryCount.ToString());
+            noPostbackContent.Headers.Add(OktaAuthenticator.TimeoutElapsedHeader, retryTimeout.ToString());
             try
             {
                 var restRequester = new Mock.MockOktaRestRequester()
@@ -41,13 +46,16 @@ public void TestMissingPostbackUrl()
                     SSOUrl = "https://snowflakecomputing.okta.com/app/snowflake_testaccountdev_1/blah/sso/saml",
                     ResponseContent = noPostbackContent,
                 };
-                var sfSession = new SFSession("account=test;user=test;password=test;authenticator=https://snowflakecomputing.okta.com;host=test", null, restRequester);
+                var sfSession = new SFSession("account=test;user=test;password=test;authenticator=https://snowflakecomputing.okta.com;" +
+                    $"host=test;MAXHTTPRETRIES={retryCount};RETRY_TIMEOUT={retryTimeout};", null, restRequester);
                 sfSession.Open();
                 Assert.Fail("Should not pass");
             } catch (SnowflakeDbException e)
             {
-                Assert.AreEqual(SFError.IDP_SAML_POSTBACK_NOTFOUND.GetAttribute<SFErrorAttr>().errorCode, e.ErrorCode);
+                Assert.AreEqual(SFError.IDP_SAML_POSTBACK_NOTFOUND.GetAttribute<SFErrorAttr>().errorCode, ((SnowflakeDbException)e.InnerException).ErrorCode);
             }
+            noPostbackContent.Headers.Remove(OktaAuthenticator.RetryCountHeader);
+            noPostbackContent.Headers.Remove(OktaAuthenticator.TimeoutElapsedHeader);
         }
 
         [Test]