Fix removeUI assertions to work with proxy UIs (#10575)

diff --git a/server/src/main/java/com/vaadin/server/VaadinSession.java b/server/src/main/java/com/vaadin/server/VaadinSession.java
index 8cb833e68f5bc91bbb31ec6f61723690a5115f3a..131f5413856f6142c134368d3bcc1b6bd0108984 100644 (file)
--- a/server/src/main/java/com/vaadin/server/VaadinSession.java
+++ b/server/src/main/java/com/vaadin/server/VaadinSession.java
@@ -867,8 +867,10 @@ public class VaadinSession implements HttpSessionBindingListener, Serializable {
      *            the UI to remove
      */
     public void removeUI(UI ui) {
-        assert hasLock();
-        assert UI.getCurrent() == ui;
+        assert hasLock() : "Session is locked";
+        assert UI.getCurrent() != null : "Current UI cannot be null";
+        assert ui != null : "Removed UI cannot be null";
+        assert UI.getCurrent().getUIId() == ui.getUIId() : "UIs don't match";
         Integer id = Integer.valueOf(ui.getUIId());
         ui.setSession(null);
         uIs.remove(id);