[flink] 03/04: [FLINK-24603][e2e] Add support for setting job main class

[ch...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

chesnay pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 295a1e352605206c9c1d91969ec3eeb9dc983bf4
Author: Chesnay Schepler <ch...@apache.org>
AuthorDate: Thu Oct 28 12:20:09 2021 +0200

    [FLINK-24603][e2e] Add support for setting job main class
---
 .../flink/tests/util/flink/FlinkDistribution.java  |  7 ++++++
 .../flink/tests/util/flink/JobSubmission.java      | 25 +++++++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/flink-end-to-end-tests/flink-end-to-end-tests-common/src/main/java/org/apache/flink/tests/util/flink/FlinkDistribution.java b/flink-end-to-end-tests/flink-end-to-end-tests-common/src/main/java/org/apache/flink/tests/util/flink/FlinkDistribution.java
index cd3d935..0988e43 100644
--- a/flink-end-to-end-tests/flink-end-to-end-tests-common/src/main/java/org/apache/flink/tests/util/flink/FlinkDistribution.java
+++ b/flink-end-to-end-tests/flink-end-to-end-tests-common/src/main/java/org/apache/flink/tests/util/flink/FlinkDistribution.java
@@ -169,6 +169,13 @@ final class FlinkDistribution {
             commands.add("-p");
             commands.add(String.valueOf(jobSubmission.getParallelism()));
         }
+        jobSubmission
+                .getMainClass()
+                .ifPresent(
+                        mainClass -> {
+                            commands.add("--class");
+                            commands.add(mainClass);
+                        });
         commands.add(jobSubmission.getJar().toAbsolutePath().toString());
         commands.addAll(jobSubmission.getArguments());
 
diff --git a/flink-end-to-end-tests/flink-end-to-end-tests-common/src/main/java/org/apache/flink/tests/util/flink/JobSubmission.java b/flink-end-to-end-tests/flink-end-to-end-tests-common/src/main/java/org/apache/flink/tests/util/flink/JobSubmission.java
index ed3eabd..90e35be 100644
--- a/flink-end-to-end-tests/flink-end-to-end-tests-common/src/main/java/org/apache/flink/tests/util/flink/JobSubmission.java
+++ b/flink-end-to-end-tests/flink-end-to-end-tests-common/src/main/java/org/apache/flink/tests/util/flink/JobSubmission.java
@@ -20,25 +20,32 @@ package org.apache.flink.tests.util.flink;
 
 import org.apache.flink.util.Preconditions;
 
+import javax.annotation.Nullable;
+
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 
 /** Programmatic definition of a job-submission. */
 public class JobSubmission {
 
     private final Path jar;
+
+    private final String mainClass;
     private final int parallelism;
     private final boolean detached;
     private final List<String> arguments;
 
     JobSubmission(
             final Path jar,
+            @Nullable final String mainClass,
             final int parallelism,
             final boolean detached,
             final List<String> arguments) {
         this.jar = jar;
+        this.mainClass = mainClass;
         this.parallelism = parallelism;
         this.detached = detached;
         this.arguments = Collections.unmodifiableList(arguments);
@@ -60,12 +67,17 @@ public class JobSubmission {
         return jar;
     }
 
+    public Optional<String> getMainClass() {
+        return Optional.ofNullable(mainClass);
+    }
+
     /** Builder for the {@link JobSubmission}. */
     public static class JobSubmissionBuilder {
         private final Path jar;
         private int parallelism = 0;
         private final List<String> arguments = new ArrayList<>(2);
         private boolean detached = false;
+        private String mainClass = null;
 
         public JobSubmissionBuilder(final Path jar) {
             Preconditions.checkNotNull(jar);
@@ -74,6 +86,17 @@ public class JobSubmission {
         }
 
         /**
+         * Sets the main class for the job.
+         *
+         * @param mainClass main class for the job
+         * @return the modified builder
+         */
+        public JobSubmissionBuilder setMainClass(final String mainClass) {
+            this.mainClass = mainClass;
+            return this;
+        }
+
+        /**
          * Sets the parallelism for the job.
          *
          * @param parallelism parallelism for the job
@@ -122,7 +145,7 @@ public class JobSubmission {
         }
 
         public JobSubmission build() {
-            return new JobSubmission(jar, parallelism, detached, arguments);
+            return new JobSubmission(jar, mainClass, parallelism, detached, arguments);
         }
     }
 }