1 files changed, 55 insertions, 9 deletions
diff --git a/server/tests/src/com/vaadin/ui/UIThemeEscaping.java b/server/tests/src/com/vaadin/ui/UIThemeEscaping.java
index ca6782952d..236f283823 100644
--- a/server/tests/src/com/vaadin/ui/UIThemeEscaping.java
+++ b/server/tests/src/com/vaadin/ui/UIThemeEscaping.java
@@ -15,29 +15,75 @@
  */
 package com.vaadin.ui;
 
-import org.junit.Assert;
+import com.vaadin.server.VaadinRequest;
+import org.junit.Before;
 import org.junit.Test;
 
-import com.vaadin.server.VaadinRequest;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.nullValue;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class UIThemeEscaping {
 
-    @Test
-    public void testThemeEscaping() {
-        UI ui = new UI() {
+    private UI ui;
+
+    private void initUiWithTheme(String theme) {
+        VaadinRequest request = getRequestWithTheme(theme);
+
+        ui.doInit(request, 1234, "foobar");
+    }
+
+    private VaadinRequest getRequestWithTheme(String theme) {
+        VaadinRequest request = mock(VaadinRequest.class);
+
+        when(request.getParameter("theme")).thenReturn(theme);
+
+        return request;
+    }
+
+    @Before
+    public void setup() {
+        ui = new UI() {
             @Override
             protected void init(VaadinRequest request) {
                 // Nothing to do
             }
         };
+    }
 
+    @Test
+    public void dangerousCharactersAreRemoved() {
         ui.setTheme("a<å(_\"$");
 
-        String theme = ui.getTheme();
+        assertThat(ui.getTheme(), is("aå_$"));
+    }
+
+    @Test
+    public void nullThemeIsSet() {
+        ui.setTheme("foobar");
+
+        ui.setTheme(null);
 
-        Assert.assertEquals(
-                "Dangerous characters should be removed from the theme name",
-                "aå_$", theme);
+        assertThat(ui.getTheme(), is(nullValue()));
     }
 
+    @Test
+    public void themeIsSetOnInit() {
+        ui.setTheme("foobar");
+
+        initUiWithTheme("bar");
+
+        assertThat(ui.getTheme(), is("bar"));
+    }
+
+    @Test
+    public void nullThemeIsSetOnInit() {
+        ui.setTheme("foobar");
+
+        initUiWithTheme(null);
+
+        assertThat(ui.getTheme(), is(nullValue()));
+    }
 }