From 20a6352fe7c232ad54a822d08ce996f785a14384 Mon Sep 17 00:00:00 2001 From: Kaituo Li Date: Thu, 8 Feb 2024 15:26:18 -0800 Subject: [PATCH] fix IT by restore env variable change Signed-off-by: Kaituo Li --- .../metrics/reporter/DimensionUtilsTest.java | 23 ++++++++++++------- .../spark/sql/util/EnvironmentProvider.scala | 10 +++++--- .../spark/sql/util/RealEnvironment.scala | 17 ++++++++++---- .../spark/sql/util/MockEnvironment.scala | 20 +++++++++++----- 4 files changed, 48 insertions(+), 22 deletions(-) diff --git a/flint-core/src/test/java/org/opensearch/flint/core/metrics/reporter/DimensionUtilsTest.java b/flint-core/src/test/java/org/opensearch/flint/core/metrics/reporter/DimensionUtilsTest.java index 7fab8c346..94760fc37 100644 --- a/flint-core/src/test/java/org/opensearch/flint/core/metrics/reporter/DimensionUtilsTest.java +++ b/flint-core/src/test/java/org/opensearch/flint/core/metrics/reporter/DimensionUtilsTest.java @@ -57,13 +57,20 @@ public void testGetDimensionsFromSystemEnv() throws NoSuchFieldException, Illega Field field = classOfMap.getDeclaredField("m"); field.setAccessible(true); Map writeableEnvironmentVariables = (Map)field.get(System.getenv()); - writeableEnvironmentVariables.put("TEST_VAR", "dummy1"); - writeableEnvironmentVariables.put("SERVERLESS_EMR_JOB_ID", "dummy2"); - Dimension result1 = DimensionUtils.constructDimension("TEST_VAR", parts); - assertEquals("TEST_VAR", result1.getName()); - assertEquals("dummy1", result1.getValue()); - Dimension result2 = DimensionUtils.constructDimension("jobId", parts); - assertEquals("jobId", result2.getName()); - assertEquals("dummy2", result2.getValue()); + try { + writeableEnvironmentVariables.put("TEST_VAR", "dummy1"); + writeableEnvironmentVariables.put("SERVERLESS_EMR_JOB_ID", "dummy2"); + Dimension result1 = DimensionUtils.constructDimension("TEST_VAR", parts); + assertEquals("TEST_VAR", result1.getName()); + assertEquals("dummy1", result1.getValue()); + Dimension result2 = DimensionUtils.constructDimension("jobId", parts); + assertEquals("jobId", result2.getName()); + assertEquals("dummy2", result2.getValue()); + } finally { + // since system environment is shared by other tests. Make sure to remove them before exiting. + writeableEnvironmentVariables.remove("SERVERLESS_EMR_JOB_ID"); + writeableEnvironmentVariables.remove("TEST_VAR"); + } + } } diff --git a/spark-sql-application/src/main/scala/org/apache/spark/sql/util/EnvironmentProvider.scala b/spark-sql-application/src/main/scala/org/apache/spark/sql/util/EnvironmentProvider.scala index e053e23fe..5b1c4e2df 100644 --- a/spark-sql-application/src/main/scala/org/apache/spark/sql/util/EnvironmentProvider.scala +++ b/spark-sql-application/src/main/scala/org/apache/spark/sql/util/EnvironmentProvider.scala @@ -9,12 +9,16 @@ package org.apache.spark.sql.util * Trait defining an interface for fetching environment variables. */ trait EnvironmentProvider { + /** * Retrieves the value of an environment variable. * - * @param name The name of the environment variable. - * @param default The default value to return if the environment variable is not set. - * @return The value of the environment variable if it exists, otherwise the default value. + * @param name + * The name of the environment variable. + * @param default + * The default value to return if the environment variable is not set. + * @return + * The value of the environment variable if it exists, otherwise the default value. */ def getEnvVar(name: String, default: String): String } diff --git a/spark-sql-application/src/main/scala/org/apache/spark/sql/util/RealEnvironment.scala b/spark-sql-application/src/main/scala/org/apache/spark/sql/util/RealEnvironment.scala index 6121c20c7..bf5eafce5 100644 --- a/spark-sql-application/src/main/scala/org/apache/spark/sql/util/RealEnvironment.scala +++ b/spark-sql-application/src/main/scala/org/apache/spark/sql/util/RealEnvironment.scala @@ -6,15 +6,22 @@ package org.apache.spark.sql.util /** - * An implementation of `EnvironmentProvider` that fetches actual environment variables from the system. + * An implementation of `EnvironmentProvider` that fetches actual environment variables from the + * system. */ class RealEnvironment extends EnvironmentProvider { + /** - * Retrieves the value of an environment variable from the system or returns a default value if not present. + * Retrieves the value of an environment variable from the system or returns a default value if + * not present. * - * @param name The name of the environment variable. - * @param default The default value to return if the environment variable is not set in the system. - * @return The value of the environment variable if it exists in the system, otherwise the default value. + * @param name + * The name of the environment variable. + * @param default + * The default value to return if the environment variable is not set in the system. + * @return + * The value of the environment variable if it exists in the system, otherwise the default + * value. */ def getEnvVar(name: String, default: String): String = sys.env.getOrElse(name, default) } diff --git a/spark-sql-application/src/test/scala/org/apache/spark/sql/util/MockEnvironment.scala b/spark-sql-application/src/test/scala/org/apache/spark/sql/util/MockEnvironment.scala index 9b9d00821..b6f3e3c97 100644 --- a/spark-sql-application/src/test/scala/org/apache/spark/sql/util/MockEnvironment.scala +++ b/spark-sql-application/src/test/scala/org/apache/spark/sql/util/MockEnvironment.scala @@ -6,17 +6,25 @@ package org.apache.spark.sql.util /** - * A mock implementation of `EnvironmentProvider` for use in tests, where environment variables can be predefined. + * A mock implementation of `EnvironmentProvider` for use in tests, where environment variables + * can be predefined. * - * @param inputMap A map representing the environment variables (name -> value). + * @param inputMap + * A map representing the environment variables (name -> value). */ class MockEnvironment(inputMap: Map[String, String]) extends EnvironmentProvider { + /** - * Retrieves the value of an environment variable from the input map or returns a default value if not present. + * Retrieves the value of an environment variable from the input map or returns a default value + * if not present. * - * @param name The name of the environment variable. - * @param default The default value to return if the environment variable is not set in the input map. - * @return The value of the environment variable from the input map if it exists, otherwise the default value. + * @param name + * The name of the environment variable. + * @param default + * The default value to return if the environment variable is not set in the input map. + * @return + * The value of the environment variable from the input map if it exists, otherwise the + * default value. */ def getEnvVar(name: String, default: String): String = inputMap.getOrElse(name, default) }