1 files changed, 20 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 3cacba1c3..ca8d56f27 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
@@ -18,6 +18,8 @@
 
 #include <gtest/gtest.h>
 
+#include <unordered_set>
+
 #include "utils/int_array_view.h"
 
 namespace latinime {
@@ -69,5 +71,23 @@ TEST(LanguageModelDictContentTest, TestUnigramProbabilityWithHistoricalInfo) {
     EXPECT_TRUE(LanguageModelDictContent.removeProbabilityEntry(wordId));
 }
 
+TEST(LanguageModelDictContentTest, TestIterateProbabilityEntry) {
+    LanguageModelDictContent languageModelDictContent(false /* useHistoricalInfo */);
+
+    const ProbabilityEntry originalEntry(0xFC, 100);
+
+    const int wordIds[] = { 1, 2, 3, 4, 5 };
+    for (const int wordId : wordIds) {
+        languageModelDictContent.setProbabilityEntry(wordId, &originalEntry);
+    }
+    std::unordered_set<int> wordIdSet(std::begin(wordIds), std::end(wordIds));
+    for (const auto entry : languageModelDictContent.getProbabilityEntries(WordIdArrayView())) {
+        EXPECT_EQ(originalEntry.getFlags(), entry.getProbabilityEntry().getFlags());
+        EXPECT_EQ(originalEntry.getProbability(), entry.getProbabilityEntry().getProbability());
+        wordIdSet.erase(entry.getWordId());
+    }
+    EXPECT_TRUE(wordIdSet.empty());
+}
+
 }  // namespace
 }  // namespace latinime