TST: Add test for draw_gantt_chart

Author: shnizzedy

nipype/pipeline/plugins/tests/test_callback.py

@@ -60,3 +60,36 @@ def test_callback_exception(tmpdir, plugin, stop_on_first_crash):
 
     sleep(0.5)  # Wait for callback to be called (python 2.7)
     assert so.statuses == [("f_node", "start"), ("f_node", "exception")]
+
+
+@pytest.mark.parametrize("plugin", ["Linear", "MultiProc", "LegacyMultiProc"])
+def test_callback_gantt(tmpdir, plugin):
+    import logging
+    import logging.handlers
+
+    from os import path
+
+    from nipype.utils.profiler import log_nodes_cb
+    from nipype.utils.draw_gantt_chart import generate_gantt_chart
+
+    log_filename = 'callback.log'
+    logger = logging.getLogger('callback')
+    logger.setLevel(logging.DEBUG)
+    handler = logging.FileHandler(log_filename)
+    logger.addHandler(handler)
+
+    #create workflow
+    wf = pe.Workflow(name="test", base_dir=tmpdir.strpath)
+    f_node = pe.Node(
+        niu.Function(function=func, input_names=[], output_names=[]), name="f_node"
+    )
+    wf.add_nodes([f_node])
+    wf.config["execution"] = {"crashdump_dir": wf.base_dir, "poll_sleep_duration": 2}
+
+    plugin_args = {"status_callback": log_nodes_cb}
+    if plugin != "Linear":
+        plugin_args['n_procs'] = 8
+    wf.run(plugin=plugin, plugin_args=plugin_args)
+
+    generate_gantt_chart('callback.log', 1 if plugin == "Linear" else 8)
+    assert path.exists('callback.log.html')

nipype/utils/draw_gantt_chart.py

@@ -153,13 +153,14 @@ def calculate_resource_timeseries(events, resource):
     all_res = 0.0
 
     # Iterate through the events
+    nan = {"Unknown", "N/A"}
     for _, event in sorted(events.items()):
         if event["event"] == "start":
-            if resource in event and event[resource] != "Unknown":
+            if resource in event and event[resource] not in nan:
                 all_res += float(event[resource])
             current_time = event["start"]
         elif event["event"] == "finish":
-            if resource in event and event[resource] != "Unknown":
+            if resource in event and event[resource] not in nan:
                 all_res -= float(event[resource])
             current_time = event["finish"]
         res[current_time] = all_res