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

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

Author: caalador

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

@@ -56,9 +56,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;
 
@@ -112,6 +110,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);
             }
         });
@@ -137,7 +138,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;
@@ -169,7 +170,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));
             }
 
@@ -226,7 +227,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));
             }
 
@@ -287,13 +288,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;
@@ -325,14 +328,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));
             }
         });
@@ -404,6 +407,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
@@ -421,24 +431,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
@@ -449,5 +460,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

@@ -90,7 +90,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);
@@ -196,7 +196,6 @@ public void handleDownloadRequest(DownloadEvent event) {
             successAtomic.set(success);
         });
 
-        customHandler.setTransferUI(ui);
         customHandler.handleDownloadRequest(downloadEvent);
 
         Assert.assertTrue(successAtomic.get());