Make VaadinServlet use enclosing UI class by default (#12039)

Change-Id: Ica2b9b9d82ef6d683cb51f3478a07e57cbfa28e5

2 files changed, 31 insertions, 0 deletions

diff --git a/server/src/com/vaadin/server/VaadinServlet.java b/server/src/com/vaadin/server/VaadinServlet.java
index a3030c9d8d..a9eb42578e 100644
--- a/server/src/com/vaadin/server/VaadinServlet.java
+++ b/server/src/com/vaadin/server/VaadinServlet.java
@@ -44,6 +44,7 @@ import com.vaadin.annotations.VaadinServletConfiguration.InitParameterName;
 import com.vaadin.sass.internal.ScssStylesheet;
 import com.vaadin.server.communication.ServletUIInitHandler;
 import com.vaadin.shared.JsonConstants;
+import com.vaadin.ui.UI;
 import com.vaadin.util.CurrentInstance;
 
 @SuppressWarnings("serial")
@@ -69,6 +70,8 @@ public class VaadinServlet extends HttpServlet implements Constants {
         super.init(servletConfig);
         Properties initParameters = new Properties();
 
+        readUiFromEnclosingClass(initParameters);
+
         readConfigurationAnnotation(initParameters);
 
         // Read default parameters from server.xml
@@ -101,6 +104,15 @@ public class VaadinServlet extends HttpServlet implements Constants {
         CurrentInstance.clearAll();
     }
 
+    private void readUiFromEnclosingClass(Properties initParameters) {
+        Class<?> enclosingClass = getClass().getEnclosingClass();
+
+        if (enclosingClass != null && UI.class.isAssignableFrom(enclosingClass)) {
+            initParameters.put(VaadinSession.UI_PARAMETER,
+                    enclosingClass.getName());
+        }
+    }
+
     private void readConfigurationAnnotation(Properties initParameters)
             throws ServletException {
         VaadinServletConfiguration configAnnotation = UIProvider
diff --git a/server/tests/src/com/vaadin/server/VaadinServletConfigurationTest.java b/server/tests/src/com/vaadin/server/VaadinServletConfigurationTest.java
index 80cb1d7b0c..d347534f5c 100644
--- a/server/tests/src/com/vaadin/server/VaadinServletConfigurationTest.java
+++ b/server/tests/src/com/vaadin/server/VaadinServletConfigurationTest.java
@@ -32,10 +32,17 @@ import org.junit.Test;
 import com.vaadin.annotations.VaadinServletConfiguration;
 import com.vaadin.server.DeploymentConfiguration.LegacyProperyToStringMode;
 import com.vaadin.server.VaadinServletConfigurationTest.MockUI;
+import com.vaadin.server.VaadinServletConfigurationTest.MockUI.ServletInUI;
 import com.vaadin.ui.UI;
 
 public class VaadinServletConfigurationTest {
     public static class MockUI extends UI {
+
+        public static class ServletInUI extends VaadinServlet {
+            // This servlet should automatically be configured to use the
+            // enclosing UI class
+        }
+
         @Override
         protected void init(VaadinRequest request) {
             // Do nothing
@@ -43,6 +50,18 @@ public class VaadinServletConfigurationTest {
     }
 
     @Test
+    public void testEnclosingUIClass() throws Exception {
+        ServletInUI servlet = new MockUI.ServletInUI();
+        servlet.init(new MockServletConfig());
+
+        Class<? extends UI> uiClass = new DefaultUIProvider()
+                .getUIClass(new UIClassSelectionEvent(new VaadinServletRequest(
+                        EasyMock.createMock(HttpServletRequest.class), servlet
+                                .getService())));
+        Assert.assertEquals(MockUI.class, uiClass);
+    }
+
+    @Test
     public void testValuesFromAnnotation() throws ServletException {
         TestServlet servlet = new TestServlet();
         servlet.init(new MockServletConfig());