From ecd9fc25ec155eeffd4ea2a75b917dd6f62dac36 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 7 Oct 2024 15:20:39 -0700 Subject: [PATCH] address review --- .../pylibcudf/tests/test_nvtext_generate_ngrams.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/python/pylibcudf/pylibcudf/tests/test_nvtext_generate_ngrams.py b/python/pylibcudf/pylibcudf/tests/test_nvtext_generate_ngrams.py index 08daba64d38..aedb4efddd0 100644 --- 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) + )