From 2708fa300a7504832d7fd155caeed6fc76d642f7 Mon Sep 17 00:00:00 2001
From: Gilberto Torrezan <gilberto@vaadin.com>
Date: Mon, 23 Jul 2018 16:14:53 +0300
Subject: [PATCH] Try not to lose the weak reference of UI during the test

---
 .../server/component/FlowClassesSerializableTest.java     | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/flow-server/src/test/java/com/vaadin/tests/server/component/FlowClassesSerializableTest.java b/flow-server/src/test/java/com/vaadin/tests/server/component/FlowClassesSerializableTest.java
index 0101b206711..018b9dcac1e 100644
--- a/flow-server/src/test/java/com/vaadin/tests/server/component/FlowClassesSerializableTest.java
+++ b/flow-server/src/test/java/com/vaadin/tests/server/component/FlowClassesSerializableTest.java
@@ -2,6 +2,7 @@
 
 import java.io.OutputStream;
 
+import org.junit.Assert;
 import org.junit.Test;
 
 import com.vaadin.flow.component.Component;
@@ -45,19 +46,18 @@ public void htmlComponentAndHtmlContainer() throws Throwable {
      */
     @Test
     public void streamResource() throws Throwable {
-        UI.setCurrent(new UI());
+        UI ui = new UI();
+        UI.setCurrent(ui);
         try {
             Element element = new Element("dummy-element");
             StreamReceiver streamReceiver = new StreamReceiver(
                     element.getNode(), "upload", new MyStreamVariable());
+            Assert.assertEquals(ui, UI.getCurrent());
             element.setAttribute("target", streamReceiver);
             serializeAndDeserialize(element);
             assertTrue("Basic smoke test with ",
                     element.getAttribute("target").length() > 10);
 
-        } catch (Throwable e) {
-            e.printStackTrace();
-            throw e;
         } finally {
             UI.setCurrent(null);
         }