From 663949da12a3ea1e44aee0b58afa5abda3ef26db Mon Sep 17 00:00:00 2001
From: Jelena Furundzic <jelena.furundzic@insightsoftware.com>
Date: Tue, 13 Aug 2024 19:22:41 -0700
Subject: [PATCH] Code review changes

---
 .../client/jdbc/SnowflakeBasicDataSource.java         |  8 ++++++--
 .../client/core/SessionUtilExternalBrowserTest.java   | 11 ++++-------
 .../snowflake/client/core/SnowflakeMFACacheTest.java  | 11 ++++-------
 3 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java
index 55c911451..d6e7b4fe0 100644
--- a/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java
+++ b/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java
@@ -388,13 +388,17 @@ public void setGetDateUseNullTimezone(Boolean getDateUseNullTimezone) {
 
   public void setEnableClientRequestMfaToken(boolean enableClientRequestMfaToken) {
     this.setAuthenticator(AUTHENTICATOR_USERNAME_PASSWORD_MFA);
-    this.properties.put("clientRequestMfaToken", enableClientRequestMfaToken);
+    this.properties.put(
+        SFSessionProperty.ENABLE_CLIENT_REQUEST_MFA_TOKEN.getPropertyKey(),
+        enableClientRequestMfaToken);
   }
 
   public void setEnableClientStoreTemporaryCredential(
       boolean enableClientStoreTemporaryCredential) {
     this.setAuthenticator(AUTHENTICATOR_EXTERNAL_BROWSER);
-    this.properties.put("clientStoreTemporaryCredential", enableClientStoreTemporaryCredential);
+    this.properties.put(
+        SFSessionProperty.ENABLE_CLIENT_STORE_TEMPORARY_CREDENTIAL.getPropertyKey(),
+        enableClientStoreTemporaryCredential);
   }
 
   public void setBrowserResponseTimeout(int seconds) {
diff --git a/src/test/java/net/snowflake/client/core/SessionUtilExternalBrowserTest.java b/src/test/java/net/snowflake/client/core/SessionUtilExternalBrowserTest.java
index 5b7ff4140..2ba00f378 100644
--- a/src/test/java/net/snowflake/client/core/SessionUtilExternalBrowserTest.java
+++ b/src/test/java/net/snowflake/client/core/SessionUtilExternalBrowserTest.java
@@ -24,7 +24,6 @@
 import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
 import java.sql.ResultSet;
-import java.sql.Statement;
 import java.util.Map;
 import net.snowflake.client.AbstractDriverIT;
 import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
@@ -259,12 +258,10 @@ public void testEnableClientStoreTemporaryCredential() throws Exception {
     ds.setUser(params.get("user"));
     ds.setEnableClientStoreTemporaryCredential(false);
 
-    try (Connection con = ds.getConnection()) {
-      Statement stmt = con.createStatement();
-      for (int i = 0; i < 3; i++) {
-        try (ResultSet rs = stmt.executeQuery("SELECT CURRENT_USER()")) {
-          assertTrue(rs.next());
-        }
+    for (int i = 0; i < 3; i++) {
+      try (Connection con = ds.getConnection();
+          ResultSet rs = con.createStatement().executeQuery("SELECT 1")) {
+        assertTrue(rs.next());
       }
     }
   }
diff --git a/src/test/java/net/snowflake/client/core/SnowflakeMFACacheTest.java b/src/test/java/net/snowflake/client/core/SnowflakeMFACacheTest.java
index 0e3792fa9..f1f0a3e73 100644
--- a/src/test/java/net/snowflake/client/core/SnowflakeMFACacheTest.java
+++ b/src/test/java/net/snowflake/client/core/SnowflakeMFACacheTest.java
@@ -22,7 +22,6 @@
 import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
-import java.sql.Statement;
 import java.util.Map;
 import java.util.Properties;
 import net.snowflake.client.AbstractDriverIT;
@@ -348,12 +347,10 @@ public void testEnableClientRequestMfaToken() throws SQLException {
     ds.setPassword(params.get("password"));
     ds.setEnableClientRequestMfaToken(false);
 
-    try (Connection con = ds.getConnection()) {
-      for (int i = 0; i < 3; i++) {
-        try (Statement stmt = con.createStatement();
-            ResultSet rs = stmt.executeQuery("SELECT CURRENT_USER()")) {
-          assertTrue(rs.next());
-        }
+    for (int i = 0; i < 3; i++) {
+      try (Connection con = ds.getConnection();
+          ResultSet rs = con.createStatement().executeQuery("SELECT 1")) {
+        assertTrue(rs.next());
       }
     }
   }