Fix Tests using Reflections

Author: ayushtkn

itests/hive-unit/src/test/java/org/apache/hive/service/auth/jwt/TestHttpJwtAuthentication.java

@@ -28,21 +28,14 @@
 import com.nimbusds.jwt.SignedJWT;
 import org.apache.hadoop.hive.conf.HiveConf;
 import org.apache.hadoop.hive.conf.HiveConf.ConfVars;
-import org.apache.hive.common.util.ReflectionUtil;
 import org.apache.hive.jdbc.HiveConnection;
-import org.apache.hive.jdbc.Utils;
 import org.apache.hive.jdbc.miniHS2.MiniHS2;
 import org.junit.AfterClass;
-import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.Test;
 
 import java.io.File;
-import java.lang.invoke.MethodHandles;
-import java.lang.invoke.VarHandle;
-import java.lang.reflect.Field;
-import java.lang.reflect.Modifier;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -53,7 +46,6 @@
 import java.sql.Statement;
 import java.util.Date;
 import java.util.HashMap;
-import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
@@ -62,15 +54,9 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
-
 public class TestHttpJwtAuthentication {
-  private static final Map<String, String> DEFAULTS = new HashMap<>(System.getenv());
-  private static Map<String, String> envMap;
-
-  private static final File jwtAuthorizedKeyFile =
-      new File("src/test/resources/auth.jwt/jwt-authorized-key.json");
-  private static final File jwtUnauthorizedKeyFile =
-      new File("src/test/resources/auth.jwt/jwt-unauthorized-key.json");
+  private static final File jwtAuthorizedKeyFile = new File("src/test/resources/auth.jwt/jwt-authorized-key.json");
+  private static final File jwtUnauthorizedKeyFile = new File("src/test/resources/auth.jwt/jwt-unauthorized-key.json");
   private static final File jwtVerificationJWKSFile =
       new File("src/test/resources/auth.jwt/jwt-verification-jwks.json");
 
@@ -82,43 +68,19 @@ public class TestHttpJwtAuthentication {
   @ClassRule
   public static final WireMockRule MOCK_JWKS_SERVER = new WireMockRule(MOCK_JWKS_SERVER_PORT);
 
-  /**
-   * This is a hack to make environment variables modifiable.
-   * Ref: https://stackoverflow.com/questions/318239/how-do-i-set-environment-variables-from-java.
-   */
-  @BeforeClass
-  public static void makeEnvModifiable() throws Exception {
-    envMap = new HashMap<>();
-    Class<?> envClass = Class.forName("java.lang.ProcessEnvironment");
-    Field theUnmodifiableEnvironmentField = envClass.getDeclaredField("theUnmodifiableEnvironment");
-    removeStaticFinalAndSetValue(theUnmodifiableEnvironmentField, envMap);
-  }
-
-  private static void removeStaticFinalAndSetValue(Field field, Object value) throws Exception {
-    ReflectionUtil.setStaticFinalFieldsModifiable(field);
-    field.set(null, value);
-  }
-
-  @Before
-  public void initEnvMap() {
-    envMap.clear();
-    envMap.putAll(DEFAULTS);
-  }
-
   @BeforeClass
   public static void setupHS2() throws Exception {
-    MOCK_JWKS_SERVER.stubFor(get("/jwks")
-        .willReturn(ok()
-            .withBody(Files.readAllBytes(jwtVerificationJWKSFile.toPath()))));
+    MOCK_JWKS_SERVER.stubFor(
+        get("/jwks").willReturn(ok().withBody(Files.readAllBytes(jwtVerificationJWKSFile.toPath()))));
 
     HiveConf conf = new HiveConf();
     conf.setBoolVar(ConfVars.HIVE_SUPPORT_CONCURRENCY, false);
     conf.setBoolVar(ConfVars.HIVE_SERVER2_LOGGING_OPERATION_ENABLED, false);
     conf.setBoolVar(ConfVars.HIVE_STATS_COL_AUTOGATHER, false);
     conf.setVar(ConfVars.HIVE_SERVER2_AUTHENTICATION, "JWT");
     // the content of the URL below is the same as jwtVerificationJWKSFile
-    conf.setVar(ConfVars.HIVE_SERVER2_AUTHENTICATION_JWT_JWKS_URL, "http://localhost:" + MOCK_JWKS_SERVER_PORT +
-        "/jwks");
+    conf.setVar(ConfVars.HIVE_SERVER2_AUTHENTICATION_JWT_JWKS_URL,
+        "http://localhost:" + MOCK_JWKS_SERVER_PORT + "/jwks");
     miniHS2 = new MiniHS2.Builder().withConf(conf).withHTTPTransport().build();
 
     miniHS2.start(new HashMap<>());
@@ -137,38 +99,38 @@ public static void stopServices() throws Exception {
   @Test
   public void testAuthorizedUser() throws Exception {
     String jwt = generateJWT(USER_1, jwtAuthorizedKeyFile.toPath(), TimeUnit.MINUTES.toMillis(5));
-    HiveConnection connection = getConnection(jwt, true);
-    assertLoggedInUser(connection, USER_1);
-    connection.close();
-
-    connection = getConnection(jwt, false);
-    assertLoggedInUser(connection, USER_1);
-    connection.close();
+    try (HiveConnection connection = getConnection(jwt)) {
+      assertLoggedInUser(connection, USER_1);
+    }
   }
 
   @Test(expected = SQLException.class)
   public void testExpiredJwt() throws Exception {
     String jwt = generateJWT(USER_1, jwtAuthorizedKeyFile.toPath(), 1);
     Thread.sleep(1);
-    HiveConnection connection = getConnection(jwt, true);
+    try (HiveConnection connection = getConnection(jwt)) {
+      // Should throw SQLException
+    }
   }
 
   @Test(expected = SQLException.class)
   public void testUnauthorizedUser() throws Exception {
     String unauthorizedJwt = generateJWT(USER_1, jwtUnauthorizedKeyFile.toPath(), TimeUnit.MINUTES.toMillis(5));
-    HiveConnection connection = getConnection(unauthorizedJwt, true);
+    try (HiveConnection connection = getConnection(unauthorizedJwt)) {
+      // Should throw SQLException
+    }
   }
 
   @Test(expected = SQLException.class)
   public void testWithoutJwtProvided() throws Exception {
-    HiveConnection connection = getConnection(null, true);
+    try (HiveConnection connection = getConnection(null)) {
+      // Should throw SQLException
+    }
   }
 
-  private HiveConnection getConnection(String jwt, Boolean putJwtInEnv) throws Exception {
+  private HiveConnection getConnection(String jwt) throws Exception {
     String url = getJwtJdbcConnectionUrl();
-    if (jwt != null && putJwtInEnv) {
-      System.getenv().put(Utils.JdbcConnectionParams.AUTH_JWT_ENV, jwt);
-    } else if (jwt != null) {
+    if (jwt != null) {
       url += "jwt=" + jwt;
     }
     Class.forName("org.apache.hive.jdbc.HiveDriver");
@@ -218,4 +180,4 @@ private void assertLoggedInUser(HiveConnection connection, String expectedUser)
     String loggedInUser = resultSet.getString(1);
     assertEquals(expectedUser, loggedInUser);
   }
-}
+}