Updated tests

planetarium · Jan 18, 2024 · 46bdf35 · 46bdf35
1 parent fa435d7
commit 46bdf35
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 32 deletions.
diff --git a/Libplanet.Store/TrieStateStore.cs b/Libplanet.Store/TrieStateStore.cs
@@ -70,8 +70,7 @@ public void PruneStates(IImmutableSet<HashDigest<SHA256>> survivingStateRootHash
             {
                 // FIXME: Bencodex fingerprints also should be tracked.
                 //        https://github.com/planetarium/libplanet/issues/1653
-                if (stateKey.Length != HashDigest<SHA256>.Size ||
-                    survivalNodes.Contains(stateKey))
+                if (survivalNodes.Contains(stateKey))
                 {
                     continue;
                 }

diff --git a/Libplanet.Tests/Store/TrieStateStoreTest.cs b/Libplanet.Tests/Store/TrieStateStoreTest.cs
@@ -69,42 +69,54 @@ public void GetStateRoot()
         [Fact]
         public void PruneStates()
         {
-            var values = ImmutableDictionary<KeyBytes, IValue>.Empty
-                .Add(new KeyBytes("foo"), (Binary)GetRandomBytes(4096))
-                .Add(
-                    new KeyBytes("bar"),
-                    (Text)ByteUtil.Hex(GetRandomBytes(2048)))
-                .Add(new KeyBytes("baz"), (Bencodex.Types.Boolean)false)
-                .Add(new KeyBytes("qux"), Bencodex.Types.Dictionary.Empty)
-                .Add(
-                    new KeyBytes("zzz"),
-                    Bencodex.Types.Dictionary.Empty
-                        .Add("binary", GetRandomBytes(4096))
-                        .Add("text", ByteUtil.Hex(GetRandomBytes(2048))));
-
             var stateStore = new TrieStateStore(_stateKeyValueStore);
-            ITrie first = stateStore.Commit(
-                values.Aggregate(
-                    stateStore.GetStateRoot(null),
-                    (prev, kv) => prev.Set(kv.Key, kv.Value)));
+            IKeyValueStore targetStateKeyValueStore = new MemoryKeyValueStore();
+            var targetStateStore = new TrieStateStore(targetStateKeyValueStore);
+            Random random = new Random();
+            List<(KeyBytes, IValue)> kvs = Enumerable.Range(0, 1_000)
+                .Select(_ =>
+                (
+                    new KeyBytes(GetRandomBytes(random.Next(20))),
+                    (IValue)new Binary(GetRandomBytes(20))
+                ))
+                .ToList();
 
-            int prevStatesCount = _stateKeyValueStore.ListKeys().Count();
-            ImmutableDictionary<KeyBytes, IValue> nextState =
-                values.SetItem(new KeyBytes("foo"), (Binary)GetRandomBytes(4096));
-            ITrie second = stateStore.Commit(
-                nextState.Aggregate(
-                    first,
-                    (prev, kv) => prev.Set(kv.Key, kv.Value)));
+            ITrie trie = stateStore.GetStateRoot(null);
+            foreach (var kv in kvs)
+            {
+                trie = trie.Set(kv.Item1, kv.Item2);
+            }
 
-            // foo = 0x666f6f
-            // updated branch node (0x6, aka root) + updated branch node (0x66) +
-            // updated short node + new value nodes
-            Assert.Equal(prevStatesCount + 4, _stateKeyValueStore.ListKeys().Count());
+            trie = stateStore.Commit(trie);
+            HashDigest<SHA256> firstHash = trie.Hash;
+            int firstStateCount = _stateKeyValueStore.ListKeys().Count();
+
+            trie = trie.Set(
+                new KeyBytes(GetRandomBytes(35)),
+                (IValue)new Binary(GetRandomBytes(20)));
+            trie = stateStore.Commit(trie);
+            HashDigest<SHA256> secondHash = trie.Hash;
+            int secondStateCount = _stateKeyValueStore.ListKeys().Count();
 
-            stateStore.PruneStates(ImmutableHashSet<HashDigest<SHA256>>.Empty.Add(second.Hash));
+            Assert.True(secondStateCount > firstStateCount);
+
+            // NOTE: Avoid possible collision of KeyBytes, just in case.
+            _stateKeyValueStore.Set(
+                new KeyBytes(GetRandomBytes(30)),
+                ByteUtil.ParseHex("00"));
+            _stateKeyValueStore.Set(
+                new KeyBytes(GetRandomBytes(40)),
+                ByteUtil.ParseHex("00"));
+            Assert.True(_stateKeyValueStore.ListKeys().Count() > secondStateCount);
 
             // It will stay at the same count of nodes.
-            Assert.Equal(prevStatesCount, _stateKeyValueStore.ListKeys().Count());
+            stateStore.PruneStates(
+                ImmutableHashSet<HashDigest<SHA256>>.Empty.Add(firstHash).Add(secondHash));
+            Assert.Equal(secondStateCount, _stateKeyValueStore.ListKeys().Count());
+
+            stateStore.PruneStates(
+                ImmutableHashSet<HashDigest<SHA256>>.Empty.Add(firstHash));
+            Assert.Equal(firstStateCount, _stateKeyValueStore.ListKeys().Count());
         }
 
         [Fact]