Added tests

snowflakedb · Oct 31, 2023 · d73b29c · d73b29c
1 parent ff49393
commit d73b29c
Show file tree

Hide file tree

Showing 3 changed files with 126 additions and 4 deletions.
diff --git a/src/test/java/net/snowflake/client/core/SessionUtilTest.java b/src/test/java/net/snowflake/client/core/SessionUtilTest.java
@@ -5,12 +5,18 @@
 package net.snowflake.client.core;
 
 import static org.hamcrest.MatcherAssert.assertThat;
-import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.*;
 
 import com.fasterxml.jackson.databind.node.BooleanNode;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import net.snowflake.client.jdbc.MockConnectionTest;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.client.utils.URIBuilder;
 import org.junit.Test;
 
 public class SessionUtilTest {
@@ -83,4 +89,42 @@ public void testConvertSystemPropertyToIntValue() {
             HttpUtil.JDBC_MAX_CONNECTIONS_PER_ROUTE_PROPERTY,
             HttpUtil.DEFAULT_MAX_CONNECTIONS_PER_ROUTE));
   }
+
+  @Test
+  public void testIsLoginRequest() {
+    List<String> testCases = new ArrayList<String>();
+    testCases.add("/session/v1/login-request");
+    testCases.add("/session/token-request");
+    testCases.add("/session/authenticator-request");
+
+    for (String testCase : testCases) {
+      try {
+        URIBuilder uriBuilder = new URIBuilder("https://test.snowflakecomputing.com");
+        uriBuilder.setPath(testCase);
+        URI uri = uriBuilder.build();
+        HttpPost postRequest = new HttpPost(uri);
+        assertTrue(SessionUtil.isLoginRequest(postRequest));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException(e);
+      }
+    }
+  }
+
+  @Test
+  public void testIsLoginRequestInvalidURIPath() {
+    List<String> testCases = new ArrayList<String>();
+    testCases.add("/session/not-a-real-path");
+
+    for (String testCase : testCases) {
+      try {
+        URIBuilder uriBuilder = new URIBuilder("https://test.snowflakecomputing.com");
+        uriBuilder.setPath(testCase);
+        URI uri = uriBuilder.build();
+        HttpPost postRequest = new HttpPost(uri);
+        assertFalse(SessionUtil.isLoginRequest(postRequest));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException(e);
+      }
+    }
+  }
 }
diff --git a/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java b/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java
@@ -554,7 +554,7 @@ public void testWrongHostNameTimeout() throws InterruptedException {
           equalTo(ErrorCode.NETWORK_ERROR.getMessageCode()));
 
       conEnd = System.currentTimeMillis();
-      assertThat("Login time out not taking effective", conEnd - connStart < 60000);
+      assertThat("Login time out not taking effective", conEnd - connStart < 300000);
 
       Thread.sleep(WAIT_FOR_TELEMETRY_REPORT_IN_MILLISECS);
       if (TelemetryService.getInstance().isDeploymentEnabled()) {
@@ -595,7 +595,7 @@ public void testHttpsLoginTimeoutWithSSL() throws InterruptedException {
           equalTo(ErrorCode.NETWORK_ERROR.getMessageCode()));
 
       conEnd = System.currentTimeMillis();
-      assertThat("Login time out not taking effective", conEnd - connStart < 60000);
+      assertThat("Login time out not taking effective", conEnd - connStart < 300000);
       Thread.sleep(WAIT_FOR_TELEMETRY_REPORT_IN_MILLISECS);
       if (TelemetryService.getInstance().isDeploymentEnabled()) {
         assertThat(

diff --git a/src/test/java/net/snowflake/client/jdbc/RestRequestTest.java b/src/test/java/net/snowflake/client/jdbc/RestRequestTest.java
@@ -29,7 +29,7 @@
 /** RestRequest unit tests. */
 public class RestRequestTest {
 
-  static final int DEFAULT_CONNECTION_TIMEOUT = 60000;
+  static final int DEFAULT_CONNECTION_TIMEOUT = 300000;
   static final int DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT = 300000; // ms
 
   private CloseableHttpResponse retryResponse() {
@@ -42,6 +42,16 @@ private CloseableHttpResponse retryResponse() {
     return retryResponse;
   }
 
+  private CloseableHttpResponse retryLoginResponse() {
+    StatusLine retryStatusLine = mock(StatusLine.class);
+    when(retryStatusLine.getStatusCode()).thenReturn(429);
+
+    CloseableHttpResponse retryResponse = mock(CloseableHttpResponse.class);
+    when(retryResponse.getStatusLine()).thenReturn(retryStatusLine);
+
+    return retryResponse;
+  }
+
   private CloseableHttpResponse successResponse() {
     StatusLine successStatusLine = mock(StatusLine.class);
     when(successStatusLine.getStatusCode()).thenReturn(200);
@@ -457,6 +467,74 @@ public CloseableHttpResponse answer(InvocationOnMock invocation) throws Throwabl
     }
   }
 
+  @Test(expected = SnowflakeSQLException.class)
+  public void testLoginMaxRetries() throws IOException, SnowflakeSQLException {
+    boolean telemetryEnabled = TelemetryService.getInstance().isEnabled();
+
+    CloseableHttpClient client = mock(CloseableHttpClient.class);
+    when(client.execute(any(HttpUriRequest.class)))
+        .thenAnswer(
+            new Answer<CloseableHttpResponse>() {
+              int callCount = 0;
+
+              @Override
+              public CloseableHttpResponse answer(InvocationOnMock invocation) throws Throwable {
+                callCount += 1;
+                if (callCount >= 4) {
+                  return retryLoginResponse();
+                } else {
+                  return socketTimeoutResponse();
+                }
+              }
+            });
+
+    try {
+      TelemetryService.disable();
+      execute(client, "/session/v1/login-request", 0, 0, 0, true, false, 1);
+      fail("testMaxRetries");
+    } finally {
+      if (telemetryEnabled) {
+        TelemetryService.enable();
+      } else {
+        TelemetryService.disable();
+      }
+    }
+  }
+
+  @Test(expected = SnowflakeSQLException.class)
+  public void testLoginTimeout() throws IOException, SnowflakeSQLException {
+    boolean telemetryEnabled = TelemetryService.getInstance().isEnabled();
+
+    CloseableHttpClient client = mock(CloseableHttpClient.class);
+    when(client.execute(any(HttpUriRequest.class)))
+        .thenAnswer(
+            new Answer<CloseableHttpResponse>() {
+              int callCount = 0;
+
+              @Override
+              public CloseableHttpResponse answer(InvocationOnMock invocation) throws Throwable {
+                callCount += 1;
+                if (callCount >= 4) {
+                  return retryLoginResponse();
+                } else {
+                  return socketTimeoutResponse();
+                }
+              }
+            });
+
+    try {
+      TelemetryService.disable();
+      execute(client, "/session/v1/login-request", 1, 0, 0, true, false, 10);
+      fail("testMaxRetries");
+    } finally {
+      if (telemetryEnabled) {
+        TelemetryService.enable();
+      } else {
+        TelemetryService.disable();
+      }
+    }
+  }
+
   @Test
   public void testMaxRetriesWithSuccessfulResponse() throws IOException {
     boolean telemetryEnabled = TelemetryService.getInstance().isEnabled();