1 files changed, 51 insertions, 35 deletions

diff --git a/icing/result/projection-tree_test.cc b/icing/result/projection-tree_test.cc
index 2b0f966..46d0c12 100644
--- a/icing/result/projection-tree_test.cc
+++ b/icing/result/projection-tree_test.cc
@@ -17,6 +17,7 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/proto/search.pb.h"
+#include "icing/schema/schema-store.h"
 
 namespace icing {
 namespace lib {
@@ -28,72 +29,87 @@ using ::testing::IsEmpty;
 using ::testing::SizeIs;
 
 TEST(ProjectionTreeTest, CreateEmptyFieldMasks) {
-  TypePropertyMask type_field_mask;
-  ProjectionTree tree(type_field_mask);
+  ProjectionTree tree({});
   EXPECT_THAT(tree.root().name, IsEmpty());
   EXPECT_THAT(tree.root().children, IsEmpty());
 }
 
 TEST(ProjectionTreeTest, CreateTreeTopLevel) {
-  TypePropertyMask type_field_mask;
-  type_field_mask.add_paths("subject");
+  SchemaStore::ExpandedTypePropertyMask type_field_mask{"", {"subject"}};
 
   ProjectionTree tree(type_field_mask);
   EXPECT_THAT(tree.root().name, IsEmpty());
   ASSERT_THAT(tree.root().children, SizeIs(1));
-  ASSERT_THAT(tree.root().children.at(0).name, Eq("subject"));
-  ASSERT_THAT(tree.root().children.at(0).children, IsEmpty());
+  EXPECT_THAT(tree.root().children.at(0).name, Eq("subject"));
+  EXPECT_THAT(tree.root().children.at(0).children, IsEmpty());
 }
 
 TEST(ProjectionTreeTest, CreateTreeMultipleTopLevel) {
-  TypePropertyMask type_field_mask;
-  type_field_mask.add_paths("subject");
-  type_field_mask.add_paths("body");
+  SchemaStore::ExpandedTypePropertyMask type_field_mask{"",
+                                                        {"subject", "body"}};
 
   ProjectionTree tree(type_field_mask);
   EXPECT_THAT(tree.root().name, IsEmpty());
   ASSERT_THAT(tree.root().children, SizeIs(2));
-  ASSERT_THAT(tree.root().children.at(0).name, Eq("subject"));
-  ASSERT_THAT(tree.root().children.at(0).children, IsEmpty());
-  ASSERT_THAT(tree.root().children.at(1).name, Eq("body"));
-  ASSERT_THAT(tree.root().children.at(1).children, IsEmpty());
+
+  const ProjectionTree::Node* child0 = &tree.root().children.at(0);
+  const ProjectionTree::Node* child1 = &tree.root().children.at(1);
+  if (child0->name != "subject") {
+    std::swap(child0, child1);
+  }
+
+  EXPECT_THAT(child0->name, Eq("subject"));
+  EXPECT_THAT(child0->children, IsEmpty());
+  EXPECT_THAT(child1->name, Eq("body"));
+  EXPECT_THAT(child1->children, IsEmpty());
 }
 
 TEST(ProjectionTreeTest, CreateTreeNested) {
-  TypePropertyMask type_field_mask;
-  type_field_mask.add_paths("subject.body");
-  type_field_mask.add_paths("body");
+  SchemaStore::ExpandedTypePropertyMask type_field_mask{
+      "", {"subject.body", "body"}};
 
   ProjectionTree tree(type_field_mask);
   EXPECT_THAT(tree.root().name, IsEmpty());
   ASSERT_THAT(tree.root().children, SizeIs(2));
-  ASSERT_THAT(tree.root().children.at(0).name, Eq("subject"));
-  ASSERT_THAT(tree.root().children.at(0).children, SizeIs(1));
-  ASSERT_THAT(tree.root().children.at(0).children.at(0).name, Eq("body"));
-  ASSERT_THAT(tree.root().children.at(0).children.at(0).children, IsEmpty());
-  ASSERT_THAT(tree.root().children.at(1).name, Eq("body"));
-  ASSERT_THAT(tree.root().children.at(1).children, IsEmpty());
+
+  const ProjectionTree::Node* child0 = &tree.root().children.at(0);
+  const ProjectionTree::Node* child1 = &tree.root().children.at(1);
+  if (child0->name != "subject.body") {
+    std::swap(child0, child1);
+  }
+
+  EXPECT_THAT(child0->name, Eq("subject"));
+  ASSERT_THAT(child0->children, SizeIs(1));
+  EXPECT_THAT(child0->children.at(0).name, Eq("body"));
+  EXPECT_THAT(child0->children.at(0).children, IsEmpty());
+  EXPECT_THAT(child1->name, Eq("body"));
+  EXPECT_THAT(child1->children, IsEmpty());
 }
 
 TEST(ProjectionTreeTest, CreateTreeNestedSharedNode) {
-  TypePropertyMask type_field_mask;
-  type_field_mask.add_paths("sender.name.first");
-  type_field_mask.add_paths("sender.emailAddress");
+  SchemaStore::ExpandedTypePropertyMask type_field_mask{
+      "", {"sender.name.first", "sender.emailAddress"}};
 
   ProjectionTree tree(type_field_mask);
   EXPECT_THAT(tree.root().name, IsEmpty());
   ASSERT_THAT(tree.root().children, SizeIs(1));
-  ASSERT_THAT(tree.root().children.at(0).name, Eq("sender"));
+  EXPECT_THAT(tree.root().children.at(0).name, Eq("sender"));
   ASSERT_THAT(tree.root().children.at(0).children, SizeIs(2));
-  ASSERT_THAT(tree.root().children.at(0).children.at(0).name, Eq("name"));
-  ASSERT_THAT(tree.root().children.at(0).children.at(0).children, SizeIs(1));
-  ASSERT_THAT(tree.root().children.at(0).children.at(0).children.at(0).name,
-              Eq("first"));
-  ASSERT_THAT(tree.root().children.at(0).children.at(0).children.at(0).children,
-              IsEmpty());
-  ASSERT_THAT(tree.root().children.at(0).children.at(1).name,
-              Eq("emailAddress"));
-  ASSERT_THAT(tree.root().children.at(0).children.at(1).children, IsEmpty());
+
+  const ProjectionTree::Node* child0_child0 =
+      &tree.root().children.at(0).children.at(0);
+  const ProjectionTree::Node* child0_child1 =
+      &tree.root().children.at(0).children.at(1);
+  if (child0_child0->name != "name") {
+    std::swap(child0_child0, child0_child1);
+  }
+
+  EXPECT_THAT(child0_child0->name, Eq("name"));
+  ASSERT_THAT(child0_child0->children, SizeIs(1));
+  EXPECT_THAT(child0_child0->children.at(0).name, Eq("first"));
+  EXPECT_THAT(child0_child0->children.at(0).children, IsEmpty());
+  EXPECT_THAT(child0_child1->name, Eq("emailAddress"));
+  EXPECT_THAT(child0_child1->children, IsEmpty());
 }
 
 }  // namespace