Fixes bug and slightly rationalises writer.

force_bdss/io/tests/test_workflow_writer.py

@@ -52,6 +52,7 @@ class TestWorkflowWriter(unittest.TestCase):
         fp = StringIO()
         wf = self._create_mock_workflow()
         wfwriter.write(wf, fp)
+        print(fp.getvalue())
         fp.seek(0)
         wfreader = WorkflowReader(self.mock_registry)
         wf_result = wfreader.read(fp)
@@ -73,3 +74,13 @@ class TestWorkflowWriter(unittest.TestCase):
             )
         ]
         return wf
+
+    def test_write_and_read_empty_workflow(self):
+        wf = Workflow()
+        wfwriter = WorkflowWriter()
+        fp = StringIO()
+        wfwriter.write(wf, fp)
+        fp.seek(0)
+        wfreader = WorkflowReader(self.mock_registry)
+        wf_result = wfreader.read(fp)
+        self.assertIsNone(wf_result.mco)

force_bdss/io/workflow_writer.py

@@ -19,17 +19,34 @@ class WorkflowWriter(HasStrictTraits):
         """
         data = {
             "version": "1",
-            "workflow": {}
         }
 
-        wf_data = data["workflow"]
-        wf_data["mco"] = {
-            "id": workflow.mco.bundle.id,
-            "model_data": workflow.mco.__getstate__()
+        data["workflow"] = self._workflow_data(workflow)
+        json.dump(data, f)
+
+    def _workflow_data(self, workflow):
+        workflow_data = {
+            "mco": self._mco_data(workflow.mco),
+            "kpi_calculators": [
+                self._bundle_model_data(kpic)
+                for kpic in workflow.kpi_calculators],
+            "data_sources": [
+                self._bundle_model_data(ds)
+                for ds in workflow.data_sources]
         }
 
+        return workflow_data
+
+    def _mco_data(self, mco):
+        """Extracts the data from the MCO object and returns its dictionary.
+        If the MCO is None, returns None"""
+        if mco is None:
+            return None
+
+        data = self._bundle_model_data(mco)
+
         parameters_data = []
-        for param in wf_data["mco"]["model_data"]["parameters"]:  # noqa
+        for param in data["model_data"]["parameters"]:
             parameters_data.append(
                 {
                     "id": param.factory.id,
@@ -37,24 +54,15 @@ class WorkflowWriter(HasStrictTraits):
                 }
             )
 
-        wf_data["mco"]["model_data"]["parameters"] = parameters_data  # noqa
-
-        kpic_data = []
-        for kpic in workflow.kpi_calculators:
-            kpic_data.append({
-                "id": kpic.bundle.id,
-                "model_data": kpic.__getstate__()}
-            )
+        data["model_data"]["parameters"] = parameters_data
+        return data
 
-        wf_data["kpi_calculators"] = kpic_data
-
-        ds_data = []
-        for ds in workflow.data_sources:
-            ds_data.append({
-                "id": ds.bundle.id,
-                "model_data": ds.__getstate__()
-            })
-
-        wf_data["data_sources"] = ds_data
+    def _bundle_model_data(self, bundle_model):
+        """
+        Extracts the data from a bundle model and returns its dictionary
+        """
+        return {
+            "id": bundle_model.bundle.id,
+            "model_data": bundle_model.__getstate__()
+        }
 
-        json.dump(data, f)