TreeDataProvider: Apply filters to children (#11868)

Fixes: https://github.com/vaadin/framework/issues/9933

Cherry pick from: https://github.com/vaadin/flow/pull/7317

3 files changed, 69 insertions, 11 deletions

diff --git a/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java b/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java
index 5436685e58..6fcd029003 100644
--- a/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java
+++ b/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java
@@ -172,9 +172,20 @@ public class TreeDataProvider<T>
 
     private Stream<T> getFilteredStream(Stream<T> stream,
             Optional<SerializablePredicate<T>> queryFilter) {
+        final Optional<SerializablePredicate<T>> combinedFilter;
         if (filter != null) {
-            stream = stream.filter(filter);
+            combinedFilter = Optional
+                    .of(queryFilter.map(filter::and).orElse(filter));
+        } else {
+            combinedFilter = queryFilter;
         }
-        return queryFilter.map(stream::filter).orElse(stream);
+        return combinedFilter.map(
+                f -> stream.filter(element -> flatten(element).anyMatch(f)))
+                .orElse(stream);
+    }
+
+    private Stream<T> flatten(T element) {
+        return Stream.concat(Stream.of(element), getTreeData()
+                .getChildren(element).stream().flatMap(this::flatten));
     }
 }
diff --git a/server/src/test/java/com/vaadin/data/provider/TreeDataProviderTest.java b/server/src/test/java/com/vaadin/data/provider/TreeDataProviderTest.java
index 7d5f1c0909..02c3f4508d 100644
--- a/server/src/test/java/com/vaadin/data/provider/TreeDataProviderTest.java
+++ b/server/src/test/java/com/vaadin/data/provider/TreeDataProviderTest.java
@@ -204,6 +204,51 @@ public class TreeDataProviderTest
     }
 
     @Test
+    public void filter_is_applied_to_children_provider_filter() {
+        final SerializablePredicate<String> dataProviderFilter = item -> item
+                .contains("Sub");
+        final HierarchicalQuery<String, SerializablePredicate<String>> query = new HierarchicalQuery<>(
+                null, null);
+        filter_is_applied_to_children(dataProviderFilter, query);
+    }
+
+    @Test
+    public void filter_is_applied_to_children_query_filter() {
+        final SerializablePredicate<String> dataProviderFilter = null;
+        final HierarchicalQuery<String, SerializablePredicate<String>> query = new HierarchicalQuery<>(
+                item -> item.contains("Sub"), null);
+        filter_is_applied_to_children(dataProviderFilter, query);
+    }
+
+    @Test
+    public void filter_is_applied_to_children_both_filters() {
+        final SerializablePredicate<String> dataProviderFilter = item -> item
+                .contains("Sub");
+        final HierarchicalQuery<String, SerializablePredicate<String>> query = new HierarchicalQuery<>(
+                dataProviderFilter, null);
+        filter_is_applied_to_children(dataProviderFilter, query);
+    }
+
+    private void filter_is_applied_to_children(
+            final SerializablePredicate<String> dataProviderFilter,
+            final HierarchicalQuery<String, SerializablePredicate<String>> query) {
+        final TreeData<String> stringData = new TreeData<>();
+        final String root1 = "Main";
+        final List<String> children1 = Arrays.asList("Sub1", "Sub2");
+        final String root2 = "Other";
+        final List<String> children2 = Arrays.asList("Foo1", "Foo2");
+        stringData.addRootItems(root1, root2);
+        stringData.addItems(root1, children1);
+        stringData.addItems(root2, children2);
+        final TreeDataProvider<String> provider = new TreeDataProvider<>(
+                stringData);
+        provider.setFilter(dataProviderFilter);
+        assertEquals("Unexpected amount of root items after filtering.", 1,
+                provider.getChildCount(query));
+        assertTrue(provider.fetchChildren(query).allMatch(root1::equals));
+    }
+
+    @Test
     public void setFilter() {
         getDataProvider().setFilter(item -> item.getValue().equals("Xyz")
                 || item.getValue().equals("Baz"));
@@ -214,7 +259,7 @@ public class TreeDataProviderTest
                 && !item.getValue().equals("Xyz"));
 
         assertEquals(
-                "Previous filter should be replaced when setting a new one", 6,
+                "Previous filter should be replaced when setting a new one", 14,
                 sizeWithUnfilteredQuery());
 
         getDataProvider().setFilter(null);
@@ -227,7 +272,7 @@ public class TreeDataProviderTest
     public void addFilter() {
         getDataProvider().addFilter(item -> item.getId() <= 10);
         getDataProvider().addFilter(item -> item.getId() >= 5);
-        assertEquals(5, sizeWithUnfilteredQuery());
+        assertEquals(8, sizeWithUnfilteredQuery());
     }
 
     @Override
@@ -240,7 +285,7 @@ public class TreeDataProviderTest
                         .size(new HierarchicalQuery<>("Xyz", null)));
         assertEquals("No item should match 'Zyx'", 0, strFilterDataProvider
                 .size(new HierarchicalQuery<>("Zyx", null)));
-        assertEquals("Unexpected number of matches for 'Foo'", 3,
+        assertEquals("Unexpected number of matches for 'Foo'", 4,
                 strFilterDataProvider
                         .size(new HierarchicalQuery<>("Foo", null)));
         assertEquals("No items should've been filtered out", rootData.size(),
@@ -256,7 +301,7 @@ public class TreeDataProviderTest
         assertEquals("No item should match 'Zyx'", 0,
                 dataProvider.size(new HierarchicalQuery<>(
                         strBean -> strBean.getValue().contains("Zyx"), null)));
-        assertEquals("Unexpected number of matches for 'Foo'", 3,
+        assertEquals("Unexpected number of matches for 'Foo'", 4,
                 getDataProvider()
                         .size(new HierarchicalQuery<>(fooFilter, null)));
     }
diff --git a/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java
index e056b99c21..31c8cdc929 100644
--- a/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java
+++ b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java
@@ -8,6 +8,7 @@ import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
 import org.junit.Before;
@@ -195,11 +196,12 @@ public class HierarchyMapperWithDataTest {
         expand(expandedNode);
 
         SerializablePredicate<Node> filter = n -> n.getNumber() % 2 == 0;
-        List<Node> expectedResult = testData.stream().filter(filter)
-                .filter(n -> roots.contains(n)
-                        || n.getParent().equals(testData.get(0))
-                        || n.getParent().equals(expandedNode))
-                .collect(Collectors.toList());
+
+        // Root nodes plus children of expanded nodes 0 and 4 that match the
+        // filter
+        List<Node> expectedResult = IntStream
+            .of(0, 1, 4, 6, 7, 10, 13, 26, 39, 52).mapToObj(testData::get)
+            .collect(Collectors.toList());
 
         mapper.setFilter(filter);