#7270: fix tests

twosigma · May 14, 2018 · f6f4754 · f6f4754
1 parent 9bab904
commit f6f4754
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/kernel/sparkex/src/main/java/com/twosigma/beakerx/scala/magic/command/SparkMagicCommand.java b/kernel/sparkex/src/main/java/com/twosigma/beakerx/scala/magic/command/SparkMagicCommand.java
@@ -32,6 +32,7 @@
 import org.apache.spark.SparkConf;
 import org.apache.spark.sql.SparkSession;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -46,19 +47,24 @@ public class SparkMagicCommand implements MagicCommandFunctionality {
   private SparkUI.SparkUIFactory sparkUIFactory;
   private SparkManager.SparkManagerFactory sparkManagerFactory;
   private SparkUI sparkUI;
-  private Map<String, SparkOption> sparkOptions = new HashMap<>();
+  private Map<String, SparkOption> sparkOptions;
 
   public SparkMagicCommand(KernelFunctionality kernel) {
     //constructor for reflection in LoadMagicMagicCommand
     this(kernel, new SparkUI.SparkUIFactoryImpl(), new SparkManagerImpl.SparkManagerFactoryImpl());
-    sparkOptions.put("--connect", this::connectToSparkSession);
-    sparkOptions.put("-c", this::connectToSparkSession);
   }
 
   SparkMagicCommand(KernelFunctionality kernel, SparkUI.SparkUIFactory sparkUIFactory, SparkManager.SparkManagerFactory sparkManagerFactory) {
     this.kernel = kernel;
     this.sparkUIFactory = sparkUIFactory;
     this.sparkManagerFactory = sparkManagerFactory;
+    configureOptions();
+  }
+
+  private void configureOptions() {
+    this.sparkOptions = new HashMap<>();
+    this.sparkOptions.put("--connect", this::connectToSparkSession);
+    this.sparkOptions.put("-c", this::connectToSparkSession);
   }
 
   @Override
@@ -95,7 +101,7 @@ private MagicCommandOutcomeItem validateOptions(List<String> options) {
   private List<String> getOptions(MagicCommandExecutionParam param) {
     String[] parts = param.getCommand().split(" ");
     if (parts.length == 1) {
-      return asList(parts);
+      return new ArrayList<>();
     }
     return asList(copyOfRange(parts, 1, parts.length));
   }

diff --git a/.../test/java/com/twosigma/beakerx/scala/magic/command/SparkMagicCommandAutoConnectTest.java b/.../test/java/com/twosigma/beakerx/scala/magic/command/SparkMagicCommandAutoConnectTest.java
@@ -49,7 +49,7 @@ public void unknownOption() {
     MagicCommandOutcomeItem execute = createSparkUi("--unknownOption");
     //then
     assertThat(execute.getStatus()).isEqualTo(MagicCommandOutcomeItem.Status.ERROR);
-    assertThat((String)execute.getMIMEContainer().get().getData()).contains("Unknown option --connect");
+    assertThat((String) execute.getMIMEContainer().get().getData()).contains("Unknown option --unknownOption");
   }
 
   @Test
@@ -74,7 +74,7 @@ public void autoConnectToSpark_by_c_option() {
 
   private MagicCommandOutcomeItem createSparkUi(String option) {
     Code code = Code.createCode("%%spark " + option, new ArrayList<>(), new ArrayList<>(), new Message());
-    MagicCommandExecutionParam param = new MagicCommandExecutionParam("%%spark --connect", "", 1, code, true);
+    MagicCommandExecutionParam param = new MagicCommandExecutionParam("%%spark " + option, "", 1, code, true);
     MagicCommandOutcomeItem execute = sparkMagicCommand.execute(param);
     return execute;
   }

diff --git a/...sparkex/src/test/java/com/twosigma/beakerx/scala/magic/command/SparkMagicCommandTest.java b/...sparkex/src/test/java/com/twosigma/beakerx/scala/magic/command/SparkMagicCommandTest.java
@@ -74,8 +74,8 @@ private MagicCommandOutcomeItem connectToSparkSecondTime() {
   }
 
   private MagicCommandOutcomeItem createSparkUiAndConnectToSession() {
-    Code code = Code.createCode("", new ArrayList<>(), new ArrayList<>(), new Message());
-    MagicCommandExecutionParam param = new MagicCommandExecutionParam("", "", 1, code, true);
+    Code code = Code.createCode("%%spark", new ArrayList<>(), new ArrayList<>(), new Message());
+    MagicCommandExecutionParam param = new MagicCommandExecutionParam("%%spark", "", 1, code, true);
     MagicCommandOutcomeItem execute = sparkMagicCommand.execute(param);
     assertThat(execute.getStatus()).isEqualTo(MagicCommandOutcomeItem.Status.OK);
     assertThat(sparkUI.isSparkSessionIsActive()).isFalse();