1 files changed, 45 insertions, 5 deletions
diff --git a/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc b/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
index 747c66b..3253509 100644
--- a/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
+++ b/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
@@ -18,6 +18,7 @@
 
 #include <numeric>
 #include <random>
+#include <unordered_set>
 
 #include "Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h"
 #include "Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h"
@@ -36,16 +37,24 @@ using SizeType = SortedMapBase::size_type;
 
 template <typename MapType>
 struct TestPolicy {
-  // TODO(mcg): increase beyond what ArraySortedMap supports
-  static const SizeType kFixedSize = SortedMapBase::kFixedSize;
+  static const SizeType kLargeSize = 100;
+};
+
+template <>
+struct TestPolicy<impl::ArraySortedMap<int, int>> {
+  // ArraySortedMap cannot insert more than this number
+  static const SizeType kLargeSize = SortedMapBase::kFixedSize;
 };
 
 template <typename IntMap>
 class SortedMapTest : public ::testing::Test {
  public:
-  template <typename Integer = SizeType>
-  Integer fixed_size() {
-    return static_cast<Integer>(TestPolicy<IntMap>::kFixedSize);
+  SortedMapBase::size_type large_size() const {
+    return TestPolicy<IntMap>::kLargeSize;
+  }
+
+  int large_number() const {
+    return static_cast<int>(large_size());
   }
 };
 
@@ -72,6 +81,37 @@ TYPED_TEST(SortedMapTest, Empty) {
   // EXPECT_TRUE(NotFound(map, 10));
 }
 
+TYPED_TEST(SortedMapTest, Size) {
+  std::mt19937 rand;
+  std::uniform_int_distribution<int> dist(0, 999);
+
+  std::unordered_set<int> expected;
+
+  TypeParam map;
+  auto n = this->large_number();
+  for (int i = 0; i < n; ++i) {
+    int value = dist(rand);
+
+    // The random number sequence can generate duplicates, so the expected size
+    // won't necessarily depend upon `i`.
+    expected.insert(value);
+
+    map = map.insert(value, value);
+    EXPECT_EQ(expected.size(), map.size());
+  }
+}
+
+TYPED_TEST(SortedMapTest, Increasing) {
+  std::vector<int> to_insert = Sequence(this->large_number());
+  TypeParam map = ToMap<TypeParam>(to_insert);
+  ASSERT_EQ(this->large_size(), map.size());
+
+  for (int i : to_insert) {
+    map = map.erase(i);
+  }
+  ASSERT_EQ(0u, map.size());
+}
+
 }  // namespace immutable
 }  // namespace firestore
 }  // namespace firebase