jarek/7561: fix location of the spark preview

kernel/base/src/main/java/com/twosigma/beakerx/widget/VBox.java

@@ -15,6 +15,8 @@
  */
 package com.twosigma.beakerx.widget;
 
+import com.twosigma.beakerx.message.Message;
+
 import java.util.List;
 
 /**
@@ -30,6 +32,12 @@ public VBox(List<Widget> children) {
     openComm();
   }
 
+
+  public VBox(List<Widget> children, Message message) {
+    super(children, message);
+    openComm(message);
+  }
+
   @Override
   public void updateValue(Object value) {
   }

kernel/base/src/test/java/com/twosigma/beakerx/evaluator/EvaluatorResultTestWatcher.java

@@ -132,12 +132,27 @@ public static Optional<Message> waitForUpdateMessage(KernelTest kernelTest) thro
     return idleMessage;
   }
 
+  public static Optional<Message> waitForDisplayDataMessage(KernelTest kernelTest) throws InterruptedException {
+    int count = 0;
+    Optional<Message> sentMessage = getDisplayDataMessage(kernelTest);
+    while (!sentMessage.isPresent() && count < ATTEMPT) {
+      Thread.sleep(SLEEP_IN_MILLIS);
+      sentMessage = getDisplayDataMessage(kernelTest);
+      count++;
+    }
+    return sentMessage;
+  }
+
+  public static Optional<Message> getDisplayDataMessage(KernelTest kernelTest) {
+    List<Message> listMessagesByType = SearchMessages.getListMessagesByType(kernelTest.getPublishedMessages(), JupyterMessages.DISPLAY_DATA);
+    return listMessagesByType.stream().findFirst();
+  }
+
   public static Optional<Message> getStreamMessage(KernelTest kernelTest) {
     List<Message> listMessagesByType = SearchMessages.getListMessagesByType(kernelTest.getPublishedMessages(), JupyterMessages.STREAM);
     return listMessagesByType.stream().findFirst();
   }
 
-
   private static Optional<Message> getIdleMessage(List<Message> messages) {
     return messages.stream().
             filter(x -> (x.type().equals(JupyterMessages.STATUS)) && (x.getContent().get("execution_state").equals("idle"))).findFirst();

kernel/sparkex/src/main/java/com/twosigma/beakerx/widget/PreviewTableDisplay.java

@@ -15,6 +15,7 @@
  */
 package com.twosigma.beakerx.widget;
 
+import com.twosigma.beakerx.message.Message;
 import com.twosigma.beakerx.table.TableDisplay;
 
 import java.util.Collection;
@@ -32,15 +33,18 @@ public class PreviewTableDisplay {
   private List<Widget> previewContent;
   private List<Widget> rowsContent;
 
-  public PreviewTableDisplay(Collection<Map<String, Object>> previewRows, Rows allRows, Count count) {
+  public PreviewTableDisplay(Collection<Map<String, Object>> previewRows, Rows allRows, Button button) {
     this.allRows = allRows;
     this.preview = previewRows;
-    this.previewContent = asList(createShowRowsButton(), new TableDisplay(this.preview));
+    this.previewContent = asList(configureShowRowsButton(button), new TableDisplay(this.preview));
     this.panel = new VBox(previewContent);
   }
 
-  private Button createShowRowsButton() {
-    Button button = new Button();
+  public PreviewTableDisplay(Collection<Map<String, Object>> previewRows, Rows allRows) {
+    this(previewRows, allRows, new Button());
+  }
+
+  private Button configureShowRowsButton(Button button) {
     button.setDescription("Preview " + ROWS + " Rows");
     button.registerOnClick((content, message) -> {
       if (this.rowsContent == null) {
@@ -50,15 +54,15 @@ private Button createShowRowsButton() {
         tableDisplay.setRowLimitMsg(String.format("Note: materializing a %s row preview of a Spark RDD", ROWS));
         this.rowsContent = asList(tableDisplay);
       }
-      changeContent(this.rowsContent);
+      changeContent(this.rowsContent, message);
     });
     return button;
   }
 
-  private void changeContent(List<Widget> content) {
+  private void changeContent(List<Widget> content, Message message) {
     panel.getLayout().setDisplayNone();
     panel.close();
-    panel = new VBox(content);
+    panel = new VBox(content, message);
     panel.display();
   }
 
@@ -69,9 +73,4 @@ public void display() {
   public interface Rows {
     Map<String, Object>[] get(int rows);
   }
-
-  public interface Count {
-    Long get();
-  }
-
-}
+}

kernel/sparkex/src/main/scala/com.twosigma.beakerx.scala.spark/SparkDisplayers.scala

@@ -22,7 +22,7 @@ import jupyter.{Displayer, Displayers}
 import java.util.{HashMap, Map}
 
 import com.twosigma.beakerx.widget.PreviewTableDisplay
-import com.twosigma.beakerx.widget.PreviewTableDisplay.{Count, Rows}
+import com.twosigma.beakerx.widget.PreviewTableDisplay.{Rows}
 
 import scala.collection.immutable.ListMap
 
@@ -45,9 +45,6 @@ object SparkDisplayers {
       preview,
       new Rows {
         override def get(rows: Int): Array[util.Map[String, AnyRef]] = com.twosigma.beakerx.scala.table.TableDisplay.toJavaMap(takeRows(ds, rows))
-      },
-      new Count {
-        override def get(): lang.Long = ds.count()
       }
     )
     previewWidget.display()

kernel/sparkex/src/test/java/com/twosigma/beakerx/widget/PreviewTableDisplayTest.java

@@ -0,0 +1,64 @@
+
+/*
+ *  Copyright 2018 TWO SIGMA OPEN SOURCE, LLC
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *         http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package com.twosigma.beakerx.widget;
+
+import com.twosigma.beakerx.KernelTest;
+import com.twosigma.beakerx.evaluator.EvaluatorResultTestWatcher;
+import com.twosigma.beakerx.kernel.KernelManager;
+import com.twosigma.beakerx.message.Message;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+import static com.twosigma.beakerx.MessageFactorTest.commMsg;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class PreviewTableDisplayTest {
+
+  private Button button;
+  private KernelTest kernel;
+
+  @Before
+  public void setUp() {
+    kernel = new KernelTest();
+    KernelManager.register(kernel);
+    this.button = new Button();
+    new PreviewTableDisplay(new ArrayList<>(), rows -> new Map[0], this.button);
+  }
+
+  @After
+  public void tearDown() {
+    KernelManager.register(null);
+  }
+
+  @Test
+  public void previewShouldHaveParentHeaderFromParentMessage() throws Exception {
+    //given
+    Message parentMessage = commMsg();
+    kernel.clearMessages();
+    //when
+    button.onClick(new HashMap(), parentMessage);
+    //then
+    Optional<Message> message = EvaluatorResultTestWatcher.waitForDisplayDataMessage(kernel);
+    assertThat(message.get().getParentHeader()).isEqualTo(parentMessage.getHeader());
+  }
+}