diff --git a/kernel/base/src/main/java/com/twosigma/beakerx/widget/VBox.java b/kernel/base/src/main/java/com/twosigma/beakerx/widget/VBox.java index fe22fff920..5868917450 100644 --- a/kernel/base/src/main/java/com/twosigma/beakerx/widget/VBox.java +++ b/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 children) { openComm(); } + + public VBox(List children, Message message) { + super(children, message); + openComm(message); + } + @Override public void updateValue(Object value) { } diff --git a/kernel/base/src/test/java/com/twosigma/beakerx/evaluator/EvaluatorResultTestWatcher.java b/kernel/base/src/test/java/com/twosigma/beakerx/evaluator/EvaluatorResultTestWatcher.java index 5cc6360a34..d3c73af6d9 100644 --- a/kernel/base/src/test/java/com/twosigma/beakerx/evaluator/EvaluatorResultTestWatcher.java +++ b/kernel/base/src/test/java/com/twosigma/beakerx/evaluator/EvaluatorResultTestWatcher.java @@ -132,12 +132,27 @@ public static Optional waitForUpdateMessage(KernelTest kernelTest) thro return idleMessage; } + public static Optional waitForDisplayDataMessage(KernelTest kernelTest) throws InterruptedException { + int count = 0; + Optional sentMessage = getDisplayDataMessage(kernelTest); + while (!sentMessage.isPresent() && count < ATTEMPT) { + Thread.sleep(SLEEP_IN_MILLIS); + sentMessage = getDisplayDataMessage(kernelTest); + count++; + } + return sentMessage; + } + + public static Optional getDisplayDataMessage(KernelTest kernelTest) { + List listMessagesByType = SearchMessages.getListMessagesByType(kernelTest.getPublishedMessages(), JupyterMessages.DISPLAY_DATA); + return listMessagesByType.stream().findFirst(); + } + public static Optional getStreamMessage(KernelTest kernelTest) { List listMessagesByType = SearchMessages.getListMessagesByType(kernelTest.getPublishedMessages(), JupyterMessages.STREAM); return listMessagesByType.stream().findFirst(); } - private static Optional getIdleMessage(List messages) { return messages.stream(). filter(x -> (x.type().equals(JupyterMessages.STATUS)) && (x.getContent().get("execution_state").equals("idle"))).findFirst(); diff --git a/kernel/sparkex/src/main/java/com/twosigma/beakerx/widget/PreviewTableDisplay.java b/kernel/sparkex/src/main/java/com/twosigma/beakerx/widget/PreviewTableDisplay.java index a3ca467b36..5e3c424541 100644 --- a/kernel/sparkex/src/main/java/com/twosigma/beakerx/widget/PreviewTableDisplay.java +++ b/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 previewContent; private List rowsContent; - public PreviewTableDisplay(Collection> previewRows, Rows allRows, Count count) { + public PreviewTableDisplay(Collection> 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> 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 content) { + private void changeContent(List 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[] get(int rows); } - - public interface Count { - Long get(); - } - -} +} \ No newline at end of file diff --git a/kernel/sparkex/src/main/scala/com.twosigma.beakerx.scala.spark/SparkDisplayers.scala b/kernel/sparkex/src/main/scala/com.twosigma.beakerx.scala.spark/SparkDisplayers.scala index af8f0110b6..a2207741c5 100644 --- a/kernel/sparkex/src/main/scala/com.twosigma.beakerx.scala.spark/SparkDisplayers.scala +++ b/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() diff --git a/kernel/sparkex/src/test/java/com/twosigma/beakerx/widget/PreviewTableDisplayTest.java b/kernel/sparkex/src/test/java/com/twosigma/beakerx/widget/PreviewTableDisplayTest.java new file mode 100644 index 0000000000..6835e198d3 --- /dev/null +++ b/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 = EvaluatorResultTestWatcher.waitForDisplayDataMessage(kernel); + assertThat(message.get().getParentHeader()).isEqualTo(parentMessage.getHeader()); + } +} \ No newline at end of file