fix: get ui from context if not given (#22833) (#22834)

fix issue where extensions are not
setting the UI to the progress
handler.

Co-authored-by: caalador <mikael.grankvist@vaadin.com>

Author: vaadin-bot

flow-server/src/main/java/com/vaadin/flow/server/streams/TransferProgressAwareHandler.java

@@ -57,9 +57,7 @@ public abstract class TransferProgressAwareHandler<T, R extends TransferProgress
     protected abstract TransferContext getTransferContext(T transferEvent);
 
     /**
-     * UI that the transfer was started with. Set from onStart(context).
-     * <p>
-     * protected for tests to set the ui.
+     * UI that the transfer was started with. Set from context if not given.
      */
     private UI ui;
 
@@ -113,6 +111,9 @@ public R whenStart(SerializableRunnable startHandler) {
         addTransferProgressListenerInternal(new TransferProgressListener() {
             @Override
             public void onStart(TransferContext context) {
+                if (ui == null) {
+                    setTransferUI(context.getUI());
+                }
                 ui.access(startHandler::run);
             }
         });
@@ -138,7 +139,7 @@ public R whenStart(SerializableConsumer<TransferContext> startHandler) {
         addTransferProgressListenerInternal(new TransferProgressListener() {
             @Override
             public void onStart(TransferContext context) {
-                ui.access(() -> startHandler.accept(context));
+                uiAccess(context, () -> startHandler.accept(context));
             }
         });
         return (R) this;
@@ -170,7 +171,7 @@ public R onProgress(
             @Override
             public void onProgress(TransferContext context,
                     long transferredBytes, long totalBytes) {
-                ui.access(() -> progressHandler.accept(context,
+                uiAccess(context, () -> progressHandler.accept(context,
                         transferredBytes, totalBytes));
             }
 
@@ -227,7 +228,7 @@ public R onProgress(SerializableBiConsumer<Long, Long> progressHandler,
             @Override
             public void onProgress(TransferContext context,
                     long transferredBytes, long totalBytes) {
-                ui.access(() -> progressHandler.accept(transferredBytes,
+                uiAccess(context, () -> progressHandler.accept(transferredBytes,
                         totalBytes));
             }
 
@@ -288,13 +289,15 @@ public R whenComplete(
         addTransferProgressListenerInternal(new TransferProgressListener() {
             @Override
             public void onError(TransferContext context, IOException reason) {
-                ui.access(() -> completeOrTerminateHandler.accept(false));
+                uiAccess(context,
+                        () -> completeOrTerminateHandler.accept(false));
             }
 
             @Override
             public void onComplete(TransferContext context,
                     long transferredBytes) {
-                ui.access(() -> completeOrTerminateHandler.accept(true));
+                uiAccess(context,
+                        () -> completeOrTerminateHandler.accept(true));
             }
         });
         return (R) this;
@@ -326,14 +329,14 @@ public R whenComplete(
         addTransferProgressListenerInternal(new TransferProgressListener() {
             @Override
             public void onError(TransferContext context, IOException reason) {
-                ui.access(() -> completeOrTerminateHandler.accept(context,
-                        false));
+                uiAccess(context, () -> completeOrTerminateHandler
+                        .accept(context, false));
             }
 
             @Override
             public void onComplete(TransferContext context,
                     long transferredBytes) {
-                ui.access(
+                uiAccess(context,
                         () -> completeOrTerminateHandler.accept(context, true));
             }
         });
@@ -405,6 +408,13 @@ protected void setTransferUI(UI ui) {
         }
     }
 
+    protected void uiAccess(TransferContext context, final Command command) {
+        if (ui == null) {
+            setTransferUI(context.getUI());
+        }
+        ui.access(command);
+    }
+
     /**
      * A wrapper for {@link TransferProgressListener} that ensures that UI
      * updates in transfer progress listeners are pushed to client
@@ -422,24 +432,25 @@ public TransferProgressListenerWrapper(
 
         @Override
         public void onStart(TransferContext context) {
-            ui.access(() -> delegate.onStart(context));
+            uiAccess(context, () -> delegate.onStart(context));
         }
 
         @Override
         public void onProgress(TransferContext context, long transferredBytes,
                 long totalBytes) {
-            ui.access(() -> delegate.onProgress(context, transferredBytes,
-                    totalBytes));
+            uiAccess(context, () -> delegate.onProgress(context,
+                    transferredBytes, totalBytes));
         }
 
         @Override
         public void onError(TransferContext context, IOException reason) {
-            ui.access(() -> delegate.onError(context, reason));
+            uiAccess(context, () -> delegate.onError(context, reason));
         }
 
         @Override
         public void onComplete(TransferContext context, long transferredBytes) {
-            ui.access(() -> delegate.onComplete(context, transferredBytes));
+            uiAccess(context,
+                    () -> delegate.onComplete(context, transferredBytes));
         }
 
         @Override
@@ -450,5 +461,13 @@ public long progressReportInterval() {
         protected void setTransferUI(UI ui) {
             this.ui = ui;
         }
+
+        protected void uiAccess(TransferContext context,
+                final Command command) {
+            if (ui == null) {
+                setTransferUI(context.getUI());
+            }
+            ui.access(command);
+        }
     }
 }

flow-server/src/test/java/com/vaadin/flow/server/streams/AbstractDownloadHandlerTest.java

@@ -92,7 +92,7 @@ public void setUp() throws IOException {
             public void handleDownloadRequest(DownloadEvent event) {
             }
         };
-        handler.setTransferUI(ui);
+
         mockContext = Mockito.mock(TransferContext.class);
         Mockito.when(mockContext.contentLength()).thenReturn(TOTAL_BYTES);
         listener = Mockito.mock(TransferProgressListener.class);
@@ -198,7 +198,6 @@ public void handleDownloadRequest(DownloadEvent event) {
             successAtomic.set(success);
         });
 
-        customHandler.setTransferUI(ui);
         customHandler.handleDownloadRequest(downloadEvent);
 
         Assert.assertTrue(successAtomic.get());