Call VaadinService::requestStart if push transport is websocket (#7489)

Fixes #7116

flow-server/src/main/java/com/vaadin/flow/server/communication/PushHandler.java

@@ -179,18 +179,16 @@ protected void suspend(AtmosphereResource resource) {
      *            the atmosphere resource for the current request
      * @param callback
      *            the push callback to call when a UI is found and locked
-     * @param websocket
-     *            true if this is a websocket message (as opposed to a HTTP
-     *            request)
      */
     private void callWithUi(final AtmosphereResource resource,
-            final PushEventCallback callback, boolean websocket) {
+            final PushEventCallback callback) {
         AtmosphereRequest req = resource.getRequest();
         VaadinServletRequest vaadinRequest = new VaadinServletRequest(req,
                 service);
         VaadinSession session = null;
 
-        if (websocket) {
+        boolean isWebsocket = resource.transport() == TRANSPORT.WEBSOCKET;
+        if (isWebsocket) {
             // For any HTTP request we have already started the request in the
             // servlet
             service.requestStart(vaadinRequest, null);
@@ -253,7 +251,7 @@ private void callWithUi(final AtmosphereResource resource,
             }
         } finally {
             try {
-                if (websocket) {
+                if (isWebsocket) {
                     service.requestEnd(vaadinRequest, null, session);
                 }
             } catch (Exception e) {
@@ -496,7 +494,7 @@ private static boolean isPushIdValid(VaadinSession session,
      *            The related atmosphere resources
      */
     void onConnect(AtmosphereResource resource) {
-        callWithUi(resource, establishCallback, false);
+        callWithUi(resource, establishCallback);
     }
 
     /**
@@ -506,8 +504,7 @@ void onConnect(AtmosphereResource resource) {
      *            The related atmosphere resources
      */
     void onMessage(AtmosphereResource resource) {
-        callWithUi(resource, receiveCallback,
-                resource.transport() == TRANSPORT.WEBSOCKET);
+        callWithUi(resource, receiveCallback);
     }
 
     /**

flow-server/src/test/java/com/vaadin/flow/server/communication/PushHandlerTest.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2000-2018 Vaadin Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.vaadin.flow.server.communication;
+
+import java.util.function.BiConsumer;
+
+import org.atmosphere.cpr.AtmosphereRequest;
+import org.atmosphere.cpr.AtmosphereResource;
+import org.atmosphere.cpr.AtmosphereResource.TRANSPORT;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import com.vaadin.flow.server.MockVaadinServletService;
+import com.vaadin.flow.server.VaadinServletService;
+
+public class PushHandlerTest {
+
+    @Test
+    public void onConnect_websocketTransport_requestStartIsCalledOnServiceInstance() {
+        VaadinServletService service = runTest((handler, resource) -> {
+            Mockito.when(resource.transport()).thenReturn(TRANSPORT.WEBSOCKET);
+            handler.onConnect(resource);
+        });
+
+        Mockito.verify(service).requestStart(Mockito.any(), Mockito.any());
+    }
+
+    @Test
+    public void onConnect_notWebsocketTransport_requestStartIsNotCalledOnServiceInstance() {
+        VaadinServletService service = runTest((handler, resource) -> {
+            Mockito.when(resource.transport()).thenReturn(TRANSPORT.AJAX);
+            handler.onConnect(resource);
+        });
+
+        Mockito.verify(service, Mockito.times(0)).requestStart(Mockito.any(),
+                Mockito.any());
+    }
+
+    @Test
+    public void onMessage_websocketTransport_requestStartIsCalledOnServiceInstance() {
+        VaadinServletService service = runTest((handler, resource) -> {
+            Mockito.when(resource.transport()).thenReturn(TRANSPORT.WEBSOCKET);
+            handler.onMessage(resource);
+        });
+
+        Mockito.verify(service).requestStart(Mockito.any(), Mockito.any());
+    }
+
+    @Test
+    public void onMessage_notWebsocketTransport_requestStartIsNotCalledOnServiceInstance() {
+        VaadinServletService service = runTest((handler, resource) -> {
+            Mockito.when(resource.transport()).thenReturn(TRANSPORT.AJAX);
+            handler.onMessage(resource);
+        });
+
+        Mockito.verify(service, Mockito.times(0)).requestStart(Mockito.any(),
+                Mockito.any());
+    }
+
+    private VaadinServletService runTest(
+            BiConsumer<PushHandler, AtmosphereResource> testExec) {
+        MockVaadinServletService service = Mockito
+                .spy(MockVaadinServletService.class);
+        service.init();
+        PushHandler handler = new PushHandler(service);
+
+        AtmosphereResource resource = Mockito.mock(AtmosphereResource.class);
+        AtmosphereRequest request = Mockito.mock(AtmosphereRequest.class);
+        Mockito.when(resource.getRequest()).thenReturn(request);
+
+        testExec.accept(handler, resource);
+
+        return service;
+    }
+}