From 3208cfc167351bd2e072040399b36c59e77d1af5 Mon Sep 17 00:00:00 2001
From: Qing Wang <kingchin1218@gmail.com>
Date: Fri, 13 May 2022 09:46:20 +0800
Subject: [PATCH] [Runtime env][Java] Add unit tests for specifying jars for
 tasks. (#24712)

It seems that we have already supported specifying java jars for normal tasks, this PR only needs to add unit tests for that.
---
 .../main/java/io/ray/test/RuntimeEnvTest.java | 53 +++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/java/test/src/main/java/io/ray/test/RuntimeEnvTest.java b/java/test/src/main/java/io/ray/test/RuntimeEnvTest.java
index a5e7c908b92ac..13f366010e8be 100644
--- a/java/test/src/main/java/io/ray/test/RuntimeEnvTest.java
+++ b/java/test/src/main/java/io/ray/test/RuntimeEnvTest.java
@@ -5,6 +5,7 @@
 import io.ray.api.Ray;
 import io.ray.api.runtimeenv.RuntimeEnv;
 import io.ray.runtime.util.SystemUtil;
+import java.util.List;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
@@ -193,4 +194,56 @@ public void testZipPackageInActor() {
     testDownloadAndLoadPackage(
         "https://github.com/ray-project/test_packages/raw/main/raw_resources/java-1.0-SNAPSHOT.zip");
   }
+
+  private static boolean findClasses(List<String> classNames) {
+    try {
+      for (String name : classNames) {
+        Class.forName(name);
+      }
+    } catch (ClassNotFoundException e) {
+      return false;
+    }
+    return true;
+  }
+
+  private static void testDownloadAndLoadPackagesForTask(
+      List<String> urls, List<String> classNames) {
+    try {
+      Ray.init();
+      final RuntimeEnv runtimeEnv = new RuntimeEnv.Builder().addJars(urls).build();
+      boolean ret =
+          Ray.task(RuntimeEnvTest::findClasses, classNames)
+              .setRuntimeEnv(runtimeEnv)
+              .remote()
+              .get();
+      Assert.assertTrue(ret);
+    } finally {
+      Ray.shutdown();
+    }
+  }
+
+  private static void testDownloadAndLoadPackagesForTask(String url, String className) {
+    testDownloadAndLoadPackagesForTask(ImmutableList.of(url), ImmutableList.of(className));
+  }
+
+  public void testJarPackageForTask() {
+    testDownloadAndLoadPackagesForTask(
+        "https://github.com/ray-project/test_packages/raw/main/raw_resources/bar.jar",
+        "io.testpackages.Bar");
+  }
+
+  public void testZipPackageForTask() {
+    testDownloadAndLoadPackagesForTask(
+        "https://github.com/ray-project/test_packages/raw/main/raw_resources/foo.zip",
+        "io.testpackages.Foo");
+  }
+
+  /// This case tests that a task needs 2 jars for load different classes.
+  public void testMultipleJars() {
+    testDownloadAndLoadPackagesForTask(
+        ImmutableList.of(
+            "https://github.com/ray-project/test_packages/raw/main/raw_resources/bar.jar",
+            "https://github.com/ray-project/test_packages/raw/main/raw_resources/foo.jar"),
+        ImmutableList.of("io.testpackages.Bar", "io.testpackages.Foo"));
+  }
 }