2 files changed, 23 insertions, 2 deletions
diff --git a/server/src/main/java/com/vaadin/tokka/server/communication/data/InMemoryDataSource.java b/server/src/main/java/com/vaadin/tokka/server/communication/data/InMemoryDataSource.java
index 7b4b84e5b9..a23d337652 100644
--- a/server/src/main/java/com/vaadin/tokka/server/communication/data/InMemoryDataSource.java
+++ b/server/src/main/java/com/vaadin/tokka/server/communication/data/InMemoryDataSource.java
@@ -30,7 +30,6 @@ import java.util.stream.Stream;
  */
 public class InMemoryDataSource<T> implements DataSource<T, Comparator<T>> {
 
-    // FIXME: Missing Query object
     private Function<Object, Stream<T>> request;
 
     /**
@@ -63,8 +62,13 @@ public class InMemoryDataSource<T> implements DataSource<T, Comparator<T>> {
     }
 
     @Override
-    public DataSource<T, Comparator<T>> sortingBy(Comparator<T> sortOrder) {
+    public InMemoryDataSource<T> sortingBy(Comparator<T> sortOrder) {
         return new InMemoryDataSource<T>(q -> request.apply(q)
                 .sorted(sortOrder));
     }
+
+    public <U extends Comparable<? super U>> InMemoryDataSource<T> sortingBy(
+            Function<T, U> sortOrder) {
+        return sortingBy(Comparator.comparing(sortOrder));
+    }
 }
diff --git a/server/src/test/java/com/vaadin/tokka/data/datasource/InMemoryDataSourceTest.java b/server/src/test/java/com/vaadin/tokka/data/datasource/InMemoryDataSourceTest.java
index 2bd96e9017..545c4f834b 100644
--- a/server/src/test/java/com/vaadin/tokka/data/datasource/InMemoryDataSourceTest.java
+++ b/server/src/test/java/com/vaadin/tokka/data/datasource/InMemoryDataSourceTest.java
@@ -83,4 +83,21 @@ public class InMemoryDataSourceTest {
             }
         }
     }
+
+    @Test
+    public void testDefatulSortWithFunction() {
+        List<StrBean> list = dataSource.sortingBy(StrBean::getValue)
+                .apply(null).collect(Collectors.toList());
+
+        Assert.assertEquals("Sorted data and original data sizes don't match",
+                data.size(), list.size());
+
+        for (int i = 1; i < list.size(); ++i) {
+            StrBean prev = list.get(i - 1);
+            StrBean cur = list.get(i);
+
+            // Test default sort
+            Assert.assertTrue(prev.getValue().compareTo(cur.getValue()) <= 0);
+        }
+    }
 }