TEST: test broadcast/auxillary variables

funpack/tests/test_processing.py

@@ -535,8 +535,7 @@ def test_processData_broadcast():
     independent,1,2,3\tmyproc(broadcast_arg=[100, 200, 300])
     """).strip()
 
-    with mock.patch('funpack.util.inMainProcess', return_value=False), \
-         tempdir():
+    with tempdir():
         open('processing.tsv', 'wt').write(procfile)
         gen_test_data(3, 50, 'data.tsv')
 
@@ -549,6 +548,122 @@ def test_processData_broadcast():
 
         processing.processData(dtable)
 
-        assert (dtable[:, '10-0.0'] == dtable[:, '1-0.0'] * 100).all()
-        assert (dtable[:, '20-0.0'] == dtable[:, '2-0.0'] * 200).all()
-        assert (dtable[:, '30-0.0'] == dtable[:, '3-0.0'] * 300).all()
+        print(dtable[:, :].columns)
+
+        assert (dtable[:, '10-0.0']  == dtable[:, '1-0.0'] * 100).all()
+        assert (dtable[:, '20-0.0']  == dtable[:, '2-0.0'] * 200).all()
+        assert (dtable[:, '30-0.0']  == dtable[:, '3-0.0'] * 300).all()
+
+
+def test_processing_auxillaryVariables():
+
+    @custom.processor(auxvids=['otherv'])
+    def myproc(dtable, vids, otherv):
+        v      = vids[0]
+        s1     = dtable[:, '{}-0.0'.format(v)]
+        s2     = dtable[:, '{}-0.0'.format(otherv)]
+        s3     = s1 * s2
+        s3.name = '{}-0.0'.format(v * 10)
+        return [s3], [v * 10]
+
+    procfile = tw.dedent("""
+    Variable\tProcess
+    independent,1,2,3\tmyproc(otherv=4)
+    """).strip()
+
+    with tempdir():
+        open('processing.tsv', 'wt').write(procfile)
+        gen_test_data(4, 50, 'data.tsv')
+
+        proctable = loadtables.loadProcessingTable('processing.tsv')
+        vartable, _, cattable = gen_tables((1, 2, 3, 4))[:3]
+        dtable, _ = importing.importData('data.tsv',
+                                         vartable,
+                                         proctable,
+                                         cattable)
+
+        processing.processData(dtable)
+
+        s4 = dtable[:, '4-0.0']
+
+        assert (dtable[:, '10-0.0'] == dtable[:, '1-0.0'] * s4).all()
+        assert (dtable[:, '20-0.0'] == dtable[:, '2-0.0'] * s4).all()
+        assert (dtable[:, '30-0.0'] == dtable[:, '3-0.0'] * s4).all()
+
+
+def test_processing_auxillaryVariables_parallel():
+
+    @custom.processor(auxvids=['otherv'])
+    def myproc(dtable, vids, otherv):
+        v      = vids[0]
+        s1     = dtable[:, '{}-0.0'.format(v)]
+        s2     = dtable[:, '{}-0.0'.format(otherv)]
+        s3     = s1 * s2
+        s3.name = '{}-0.0'.format(v * 10)
+        return [s3], [v * 10]
+
+    procfile = tw.dedent("""
+    Variable\tProcess
+    independent,1,2,3\tmyproc(otherv=4)
+    """).strip()
+
+    with tempdir():
+        open('processing.tsv', 'wt').write(procfile)
+        gen_test_data(4, 50, 'data.tsv')
+
+        pool      = mp.Pool()
+        mgr       = mp.Manager()
+
+        proctable = loadtables.loadProcessingTable('processing.tsv')
+        vartable, _, cattable = gen_tables((1, 2, 3, 4))[:3]
+        dtable, _ = importing.importData('data.tsv',
+                                         vartable,
+                                         proctable,
+                                         cattable,
+                                         pool=pool,
+                                         mgr=mgr)
+
+        processing.processData(dtable)
+
+        s4 = dtable[:, '4-0.0']
+
+        assert (dtable[:, '10-0.0'] == dtable[:, '1-0.0'] * s4).all()
+        assert (dtable[:, '20-0.0'] == dtable[:, '2-0.0'] * s4).all()
+        assert (dtable[:, '30-0.0'] == dtable[:, '3-0.0'] * s4).all()
+
+
+def test_processing_broadcast_auxillaryVariables():
+    @custom.processor(auxvids=['otherv'])
+    def myproc(dtable, vids, otherv):
+        v      = vids[0]
+        s1     = dtable[:, '{}-0.0'.format(v)]
+        s2     = dtable[:, '{}-0.0'.format(otherv)]
+        s3     = s1 * s2
+        s3.name = '{}-0.0'.format(v * 10)
+        return [s3], [v * 10]
+
+    procfile = tw.dedent("""
+    Variable\tProcess
+    independent,1,2,3\tmyproc(broadcast_otherv=[4, 5, 6])
+    """).strip()
+
+    with tempdir():
+        open('processing.tsv', 'wt').write(procfile)
+        gen_test_data(6, 50, 'data.tsv')
+
+        proctable = loadtables.loadProcessingTable('processing.tsv')
+        vartable, _, cattable = gen_tables((1, 2, 3, 4, 5, 6))[:3]
+        dtable, _ = importing.importData('data.tsv',
+                                         vartable,
+                                         proctable,
+                                         cattable)
+
+        processing.processData(dtable)
+
+        s4 = dtable[:, '4-0.0']
+        s5 = dtable[:, '5-0.0']
+        s6 = dtable[:, '6-0.0']
+
+        assert (dtable[:, '10-0.0'] == dtable[:, '1-0.0'] * s4).all()
+        assert (dtable[:, '20-0.0'] == dtable[:, '2-0.0'] * s5).all()
+        assert (dtable[:, '30-0.0'] == dtable[:, '3-0.0'] * s6).all()