fix: set Vaadin session attribute using lock in reinitializeSession (#12401)

* Cherry picked unit test from Flow

See https://github.com/vaadin/flow/pull/11538

* Fix missing import

* Cherry pick fix from Flow

1 files changed, 44 insertions, 0 deletions

diff --git a/server/src/test/java/com/vaadin/server/VaadinServiceTest.java b/server/src/test/java/com/vaadin/server/VaadinServiceTest.java
index f10f55db36..f179ad1d58 100644
--- a/server/src/test/java/com/vaadin/server/VaadinServiceTest.java
+++ b/server/src/test/java/com/vaadin/server/VaadinServiceTest.java
@@ -8,8 +8,11 @@ import javax.servlet.ServletConfig;
 import javax.servlet.ServletException;
 import javax.servlet.http.HttpSessionBindingEvent;
 
+import java.util.Collections;
+
 import org.easymock.EasyMock;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import com.vaadin.shared.Registration;
 import com.vaadin.util.CurrentInstance;
@@ -179,4 +182,45 @@ public class VaadinServiceTest {
         assertEquals(1, listener.callCount);
         assertEquals(1, listener2.callCount);
     }
+
+    @Test
+    public void reinitializeSession_setVaadinSessionAttriuteWithLock() {
+        VaadinRequest request = Mockito.mock(VaadinRequest.class);
+
+        VaadinSession vaadinSession = Mockito.mock(VaadinSession.class);
+        VaadinSession newVaadinSession = Mockito.mock(VaadinSession.class);
+
+        WrappedSession session = mockSession(request, vaadinSession, "foo");
+
+        Mockito.doAnswer(invocation -> {
+            mockSession(request, newVaadinSession, "bar");
+            return null;
+        }).when(session).invalidate();
+
+        VaadinService.reinitializeSession(request);
+
+        Mockito.verify(vaadinSession, Mockito.times(2)).lock();
+        Mockito.verify(vaadinSession).setAttribute(
+                VaadinService.PRESERVE_UNBOUND_SESSION_ATTRIBUTE, Boolean.TRUE);
+        Mockito.verify(vaadinSession).setAttribute(
+                VaadinService.PRESERVE_UNBOUND_SESSION_ATTRIBUTE, null);
+        Mockito.verify(vaadinSession, Mockito.times(2)).unlock();
+    }
+
+    private WrappedSession mockSession(VaadinRequest request,
+            VaadinSession vaadinSession, String attributeName) {
+        WrappedSession session = Mockito.mock(WrappedSession.class);
+        Mockito.when(request.getWrappedSession()).thenReturn(session);
+
+        Mockito.when(session.getAttributeNames())
+                .thenReturn(Collections.singleton(attributeName));
+
+        Mockito.when(session.getAttribute(attributeName))
+                .thenReturn(vaadinSession);
+
+        VaadinService service = Mockito.mock(VaadinService.class);
+
+        Mockito.when(vaadinSession.getService()).thenReturn(service);
+        return session;
+    }
 }