feat: add test for chunking logic

Aarhus-Psychiatry-Research · Sep 30, 2022 · 199ee6b · 199ee6b
1 parent 94af649
commit 199ee6b
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/tests/test_timeseriesflattener/test_add_values.py b/tests/test_timeseriesflattener/test_add_values.py
@@ -13,6 +13,7 @@
 
 from psycop_feature_generation.loaders.raw.load_text import (  # noqa pylint: disable=unused-import
     load_synth_notes,
+    _chunk_text,
 )
 from psycop_feature_generation.timeseriesflattener import (
     FlattenedDataset,
@@ -579,3 +580,16 @@ def test_add_hf_text_data():
 
     # 768 nas = 2 ids * 384 predictors with lookbehind 1 day. First get sum of each column. Then get sum of the row.
     assert outcome_df.isna().sum().sum() == 768
+
+
+def test_chunk_text():
+    text = "This is a test. This is another test. This is a third test. This is a fourth test."
+    expected = [
+        'This is a test.',
+        'This is another test.',
+        'This is a third',
+        'test. This is a'
+        ]
+
+    assert _chunk_text(text, 4) == expected
+