2 files changed, 16 insertions, 0 deletions
diff --git a/native/jni/tests/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content_test.cpp b/native/jni/tests/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content_test.cpp
index 6eef2040b..a1a3b1a75 100644
--- a/native/jni/tests/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content_test.cpp
+++ b/native/jni/tests/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content_test.cpp
@@ -35,6 +35,13 @@ TEST(LanguageModelDictContentTest, TestUnigramProbability) {
             LanguageModelDictContent.getProbabilityEntry(wordId);
     EXPECT_EQ(flag, entry.getFlags());
     EXPECT_EQ(probability, entry.getProbability());
+
+    // Remove
+    EXPECT_TRUE(LanguageModelDictContent.removeProbabilityEntry(wordId));
+    EXPECT_FALSE(LanguageModelDictContent.getProbabilityEntry(wordId).isValid());
+    EXPECT_FALSE(LanguageModelDictContent.removeProbabilityEntry(wordId));
+    EXPECT_TRUE(LanguageModelDictContent.setProbabilityEntry(wordId, &probabilityEntry));
+    EXPECT_TRUE(LanguageModelDictContent.getProbabilityEntry(wordId).isValid());
 }
 
 TEST(LanguageModelDictContentTest, TestUnigramProbabilityWithHistoricalInfo) {
@@ -53,6 +60,13 @@ TEST(LanguageModelDictContentTest, TestUnigramProbabilityWithHistoricalInfo) {
     EXPECT_EQ(timestamp, entry.getHistoricalInfo()->getTimeStamp());
     EXPECT_EQ(level, entry.getHistoricalInfo()->getLevel());
     EXPECT_EQ(count, entry.getHistoricalInfo()->getCount());
+
+    // Remove
+    EXPECT_TRUE(LanguageModelDictContent.removeProbabilityEntry(wordId));
+    EXPECT_FALSE(LanguageModelDictContent.getProbabilityEntry(wordId).isValid());
+    EXPECT_FALSE(LanguageModelDictContent.removeProbabilityEntry(wordId));
+    EXPECT_TRUE(LanguageModelDictContent.setProbabilityEntry(wordId, &probabilityEntry));
+    EXPECT_TRUE(LanguageModelDictContent.removeProbabilityEntry(wordId));
 }
 
 }  // namespace
diff --git a/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp b/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp
index 8c8e8838a..9904a3aa8 100644
--- a/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp
+++ b/native/jni/tests/suggest/policyimpl/dictionary/utils/trie_map_test.cpp
@@ -70,6 +70,8 @@ TEST(TrieMapTest, TestRemove) {
     EXPECT_FALSE(result.mIsValid);
     EXPECT_EQ(TrieMap::INVALID_INDEX, result.mNextLevelBitmapEntryIndex);
     EXPECT_EQ(11ull, trieMap.getRoot(12).mValue);
+    EXPECT_TRUE(trieMap.putRoot(S_INT_MAX, 0xFFFFFFFFFull));
+    EXPECT_TRUE(trieMap.remove(S_INT_MAX, trieMap.getRootBitmapEntryIndex()));
 }
 
 TEST(TrieMapTest, TestSetAndGetLarge) {