TEST: Unit test for flattenHierarchical

ukbparse/tests/test_cleaning.py

@@ -11,6 +11,7 @@ import string
 import textwrap as tw
 
 import multiprocessing as mp
+from unittest import mock
 
 import pytest
 
@@ -591,3 +592,74 @@ def _test_parseSpirometryData(lowMemory):
         dt = None
         mgr = None
         pool = None
+
+
+def  test_flattenHierarchical_lowMemory(): _test_flattenHierarchical(True)
+def  test_flattenHierarchical():           _test_flattenHierarchical(False)
+def _test_flattenHierarchical(lowMemory):
+
+    data = tw.dedent("""
+    eid,1-0.0
+    1,10
+    2,20
+    3,30
+    4,40
+    5,50
+    6,60
+    7,70
+    8,80
+    9,90
+    """).strip()
+
+    hier = tw.dedent("""
+    coding	meaning	node_id	parent_id
+    10	meaning 10	1	0
+    20	meaning 20	2	0
+    30	meaning 30	3	1
+    40	meaning 40	4	1
+    50	meaning 50	5	4
+    60	meaning 60	6	2
+    70	meaning 70	7	2
+    80	meaning 80	8	6
+    90	meaning 90	9	8
+    """).strip()
+
+    tests = [
+        (0, [10, 20, 10, 10, 10, 20, 20, 20, 20]),
+        (1, [10, 20, 30, 40, 40, 60, 70, 60, 60]),
+        (2, [10, 20, 30, 40, 50, 60, 70, 80, 80]),
+        (3, [10, 20, 30, 40, 50, 60, 70, 80, 90]),
+        (4, [10, 20, 30, 40, 50, 60, 70, 80, 90]),
+    ]
+
+    with tempdir(), mp.Pool(mp.cpu_count()) as pool:
+
+        with open('data.txt', 'wt') as f:
+            f.write(data)
+
+        with open('hier.txt', 'wt') as f:
+            f.write(hier)
+
+        mgr = mp.Manager()
+        vartable, proctable, cattable, _ = gen_tables([1])
+
+        for level, exp in tests:
+
+            dt, _ = importing.importData('data.txt',
+                                         vartable,
+                                         proctable,
+                                         cattable,
+                                         removeUnknown=False,
+                                         lowMemory=lowMemory,
+                                         pool=pool,
+                                         mgr=mgr)
+
+            with mock.patch('ukbparse.hierarchy.getHierarchyFilePath',
+                            return_value='hier.txt'):
+                cfns.flattenHierarchical(dt, 1, level)
+
+            assert np.all(dt[:, '1-0.0'].values == exp)
+
+        dt = None
+        mgr = None
+        pool = None