diff --git a/force_bdss/io/tests/test_workflow_writer.py b/force_bdss/io/tests/test_workflow_writer.py index df70880bf6cd64b44e8c274319eb1459ecc32b54..4ce0b2ac974af878e5369d7efbde3bdaa81c0152 100644 --- a/force_bdss/io/tests/test_workflow_writer.py +++ b/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) diff --git a/force_bdss/io/workflow_writer.py b/force_bdss/io/workflow_writer.py index 5ba53f1ab4c85bdcb8677f1076102009e9555168..79d4b6a7fc1f39bd6835a7e31187290b47ed5ce2 100644 --- a/force_bdss/io/workflow_writer.py +++ b/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)