address review

rapidsai · Oct 7, 2024 · ecd9fc2 · ecd9fc2
1 parent 577a7fc
commit ecd9fc2
Showing 1 changed file with 4 additions and 8 deletions.
diff --git a/python/pylibcudf/pylibcudf/tests/test_nvtext_generate_ngrams.py b/python/pylibcudf/pylibcudf/tests/test_nvtext_generate_ngrams.py
@@ -45,11 +45,7 @@ def test_hash_character_ngrams(input_col, ngram):
         ngram,
     )
     pa_result = plc.interop.to_arrow(result)
-    if ngram == 2:
-        assert len(pa_result[0]) == 1
-        assert len(pa_result[1]) == 2
-        assert len(pa_result[2]) == 2
-    else:
-        assert len(pa_result[0]) == 0
-        assert len(pa_result[1]) == 1
-        assert len(pa_result[2]) == 1
+    assert all(
+        len(got) == max(0, len(s.as_py()) - ngram + 1)
+        for got, s in zip(pa_result, input_col)
+    )