Added unit tests

apache · Nov 13, 2024 · 70cab90 · 70cab90
1 parent 0a98777
commit 70cab90
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 4 deletions.
diff --git a/gobblin-temporal/src/main/java/org/apache/gobblin/temporal/yarn/YarnService.java b/gobblin-temporal/src/main/java/org/apache/gobblin/temporal/yarn/YarnService.java
@@ -246,7 +246,7 @@ public YarnService(Config config, String applicationName, String applicationId,
         config.getString(GobblinYarnConfigurationKeys.YARN_APPLICATION_PROXY_JVM_ARGS) : StringUtils.EMPTY;
 
     //We get config value as emptyStringPlaceholder when the string is actually supposed to be empty
-    this.proxyJvmArgs = proxyConfigValue.equals(GobblinYarnConfigurationKeys.EMPTY_STRING_PLACEHOLDER)
+    this.proxyJvmArgs = proxyConfigValue.contains(GobblinYarnConfigurationKeys.EMPTY_STRING_PLACEHOLDER)
         ? StringUtils.EMPTY : proxyConfigValue;
 
     int numContainerLaunchThreads =

diff --git a/gobblin-yarn/src/main/java/org/apache/gobblin/yarn/GobblinYarnAppLauncher.java b/gobblin-yarn/src/main/java/org/apache/gobblin/yarn/GobblinYarnAppLauncher.java
@@ -286,7 +286,7 @@ public GobblinYarnAppLauncher(Config config, YarnConfiguration yarnConfiguration
         config.getString(GobblinYarnConfigurationKeys.YARN_APPLICATION_PROXY_JVM_ARGS) : StringUtils.EMPTY;
 
     //We get config value as emptyStringPlaceholder when the string is actually supposed to be empty
-    this.proxyJvmArgs = proxyConfigValue.equals(GobblinYarnConfigurationKeys.EMPTY_STRING_PLACEHOLDER)
+    this.proxyJvmArgs = proxyConfigValue.contains(GobblinYarnConfigurationKeys.EMPTY_STRING_PLACEHOLDER)
         ? StringUtils.EMPTY : proxyConfigValue;
 
     this.sinkLogRootDir = ConfigUtils.getString(config, GobblinYarnConfigurationKeys.LOGS_SINK_ROOT_DIR_KEY, null);

diff --git a/gobblin-yarn/src/test/java/org/apache/gobblin/yarn/GobblinYarnAppLauncherTest.java b/gobblin-yarn/src/test/java/org/apache/gobblin/yarn/GobblinYarnAppLauncherTest.java
@@ -215,7 +215,7 @@ public void setUp() throws Exception {
         .resolve();
   }
 
-  @Test
+  @Test(dependsOnMethods = "testCreateHelixCluster")
   public void testBuildApplicationMasterCommand() {
     String command = this.gobblinYarnAppLauncher.buildApplicationMasterCommand("application_1234_3456", 64);
 
@@ -355,6 +355,23 @@ public void testSendShutdownRequest() throws Exception {
     assertWithBackoff.assertEquals(getInstanceMessageNum, 0, "all controller messages processed");
   }
 
+  @Test(dependsOnMethods = "testCreateHelixCluster")
+  public void testProxyJvmArgs() throws IOException {
+    this.config = this.config.withValue(GobblinYarnConfigurationKeys.YARN_APPLICATION_PROXY_JVM_ARGS,
+        ConfigValueFactory.fromAnyRef("-Dhttp.proxyHost=zwus2-proxy.azure.linkedin.com -Dhttp.proxyPort=3128"));
+    this.gobblinYarnAppLauncher = new GobblinYarnAppLauncher(this.config, clusterConf);
+
+    String command = this.gobblinYarnAppLauncher.buildApplicationMasterCommand("application_1234_3456", 64);
+    Assert.assertTrue(command.contains("-Dhttp.proxyHost=zwus2-proxy.azure.linkedin.com -Dhttp.proxyPort=3128"));
+
+    this.config = this.config.withValue(GobblinYarnConfigurationKeys.YARN_APPLICATION_PROXY_JVM_ARGS,
+        ConfigValueFactory.fromAnyRef("emptyStringPlaceholder"));
+    this.gobblinYarnAppLauncher = new GobblinYarnAppLauncher(this.config, clusterConf);
+
+    command = this.gobblinYarnAppLauncher.buildApplicationMasterCommand("application_1234_3456", 64);
+    Assert.assertFalse(command.contains("emptyStringPlaceholder"));
+  }
+
   @AfterClass
   public void tearDown() throws IOException, TimeoutException {
     try {

diff --git a/gobblin-yarn/src/test/java/org/apache/gobblin/yarn/YarnServiceTest.java b/gobblin-yarn/src/test/java/org/apache/gobblin/yarn/YarnServiceTest.java
@@ -129,7 +129,9 @@ public void testYarnContainerStartupCommand() throws Exception{
 
     Config modifiedConfig = this.config
         .withValue(GobblinYarnConfigurationKeys.CONTAINER_JVM_MEMORY_OVERHEAD_MBS_KEY, ConfigValueFactory.fromAnyRef("10"))
-        .withValue(GobblinYarnConfigurationKeys.CONTAINER_JVM_MEMORY_XMX_RATIO_KEY, ConfigValueFactory.fromAnyRef("0.8"));
+        .withValue(GobblinYarnConfigurationKeys.CONTAINER_JVM_MEMORY_XMX_RATIO_KEY, ConfigValueFactory.fromAnyRef("0.8"))
+        .withValue(GobblinYarnConfigurationKeys.YARN_APPLICATION_PROXY_JVM_ARGS,
+        ConfigValueFactory.fromAnyRef("-Dhttp.proxyHost=zwus2-proxy.azure.linkedin.com -Dhttp.proxyPort=3128"));
     TestYarnService yarnService = new TestYarnService(modifiedConfig, "testApp2", "appId2",
             this.clusterConf, FileSystem.getLocal(new Configuration()), this.eventBus);
 
@@ -147,6 +149,7 @@ public void testYarnContainerStartupCommand() throws Exception{
     Assert.assertTrue(command.contains(String.format("--%s %s", HELIX_INSTANCE_NAME_OPTION_NAME, helixInstance)));
     Assert.assertTrue(command.contains(String.format("--%s %s", HELIX_INSTANCE_TAGS_OPTION_NAME, helixTag)));
     Assert.assertTrue(command.endsWith("1><LOG_DIR>/GobblinYarnTaskRunner.stdout 2><LOG_DIR>/GobblinYarnTaskRunner.stderr"));
+    Assert.assertFalse(command.contains("-Dhttp.proxyHost=zwus2-proxy.azure.linkedin.com -Dhttp.proxyPort=3128"));
   }
 
   static class TestYarnService extends YarnService {