2 files changed, 47 insertions, 13 deletions
diff --git a/server/tests/src/com/vaadin/tests/util/MockUI.java b/server/tests/src/com/vaadin/tests/util/MockUI.java
new file mode 100644
index 0000000000..17dc24e9e8
--- /dev/null
+++ b/server/tests/src/com/vaadin/tests/util/MockUI.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2000-2014 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.tests.util;
+
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.server.VaadinSession;
+import com.vaadin.ui.UI;
+
+public class MockUI extends UI {
+
+    public MockUI() {
+        this(findOrcreateSession());
+    }
+
+    public MockUI(VaadinSession session) {
+        setSession(session);
+        setCurrent(this);
+    }
+
+    @Override
+    protected void init(VaadinRequest request) {
+        // Do nothing
+    }
+
+    private static VaadinSession findOrcreateSession() {
+        VaadinSession session = VaadinSession.getCurrent();
+        if (session == null) {
+            session = new AlwaysLockedVaadinSession(null);
+            VaadinSession.setCurrent(session);
+        }
+        return session;
+    }
+}
diff --git a/server/tests/src/com/vaadin/ui/LabelDataSourceTest.java b/server/tests/src/com/vaadin/ui/LabelDataSourceTest.java
index 8ec3ca1245..030504cf2b 100644
--- a/server/tests/src/com/vaadin/ui/LabelDataSourceTest.java
+++ b/server/tests/src/com/vaadin/ui/LabelDataSourceTest.java
@@ -22,9 +22,9 @@ import org.junit.Before;
 import org.junit.Test;
 
 import com.vaadin.data.util.ObjectProperty;
-import com.vaadin.server.VaadinRequest;
 import com.vaadin.server.VaadinSession;
 import com.vaadin.tests.util.AlwaysLockedVaadinSession;
+import com.vaadin.tests.util.MockUI;
 
 public class LabelDataSourceTest {
 
@@ -104,18 +104,6 @@ public class LabelDataSourceTest {
         Assert.assertEquals("after", label.getValue());
     }
 
-    public class MockUI extends UI {
-
-        public MockUI() {
-            setSession(vaadinSession);
-        }
-
-        @Override
-        protected void init(VaadinRequest request) {
-        }
-
-    }
-
     @Test
     public void attachToSessionWithDifferentLocale() {
         label.setValue("before");