From 8750e85f1a5e4fd51b3ee1e477fe7ed2ab49535e Mon Sep 17 00:00:00 2001
From: Stefano Borini <sborini@enthought.com>
Date: Thu, 20 Jul 2017 11:57:08 +0100
Subject: [PATCH] Added accurate workflow writer and tests.
---
force_bdss/io/tests/test_workflow_writer.py | 65 +++++++++++++++++++++
force_bdss/io/workflow_reader.py | 8 ++-
force_bdss/io/workflow_writer.py | 30 +++++-----
force_bdss/workspecs/workflow.py | 2 +-
4 files changed, 89 insertions(+), 16 deletions(-)
create mode 100644 force_bdss/io/tests/test_workflow_writer.py
diff --git a/force_bdss/io/tests/test_workflow_writer.py b/force_bdss/io/tests/test_workflow_writer.py
new file mode 100644
index 0000000..af25d56
--- /dev/null
+++ b/force_bdss/io/tests/test_workflow_writer.py
@@ -0,0 +1,65 @@
+import unittest
+import json
+from six import StringIO
+
+from force_bdss.bundle_registry_plugin import BundleRegistryPlugin
+from force_bdss.io.workflow_reader import WorkflowReader
+
+try:
+ import mock
+except ImportError:
+ from unittest import mock
+
+from force_bdss.id_generators import bundle_id
+from force_bdss.io.workflow_writer import WorkflowWriter
+from force_bdss.mco.base_mco_model import BaseMCOModel
+from force_bdss.mco.i_multi_criteria_optimizer_bundle import \
+ IMultiCriteriaOptimizerBundle
+from force_bdss.workspecs.workflow import Workflow
+
+
+class TestWorkflowWriter(unittest.TestCase):
+ def setUp(self):
+ self.mock_registry = mock.Mock(spec=BundleRegistryPlugin)
+ mock_mco_bundle = mock.Mock(spec=IMultiCriteriaOptimizerBundle,
+ id=bundle_id("enthought", "mock"))
+ mock_mco_model = mock.Mock(
+ spec=BaseMCOModel,
+ bundle=mock_mco_bundle
+ )
+ mock_mco_bundle.create_model = mock.Mock(
+ return_value = mock_mco_model
+ )
+ self.mock_registry.mco_bundle_by_id = mock.Mock(
+ return_value=mock_mco_bundle)
+
+ def test_write(self):
+ wfwriter = WorkflowWriter()
+ fp = StringIO()
+ wf = self._create_mock_workflow()
+ wfwriter.write(wf, fp)
+ result = json.loads(fp.getvalue())
+ self.assertIn("version", result)
+ self.assertIn("workflow", result)
+ self.assertIn("multi_criteria_optimizer", result["workflow"])
+ self.assertIn("data_sources", result["workflow"])
+ self.assertIn("kpi_calculators", result["workflow"])
+
+ def test_write_and_read(self):
+ wfwriter = WorkflowWriter()
+ fp = StringIO()
+ wf = self._create_mock_workflow()
+ wfwriter.write(wf, fp)
+ fp.seek(0)
+ wfreader = WorkflowReader(self.mock_registry)
+ wf_result = wfreader.read(fp)
+ self.assertEqual(wf_result.multi_criteria_optimizer.bundle.id,
+ wf.multi_criteria_optimizer.bundle.id)
+
+ def _create_mock_workflow(self):
+ wf = Workflow()
+ wf.multi_criteria_optimizer = BaseMCOModel(
+ mock.Mock(
+ spec=IMultiCriteriaOptimizerBundle,
+ id=bundle_id("enthought", "mock")))
+ return wf
diff --git a/force_bdss/io/workflow_reader.py b/force_bdss/io/workflow_reader.py
index 643da8c..78b5a94 100644
--- a/force_bdss/io/workflow_reader.py
+++ b/force_bdss/io/workflow_reader.py
@@ -60,7 +60,13 @@ class WorkflowReader(HasStrictTraits):
def _extract_mco(self, json_data):
registry = self.bundle_registry
- mco_id = json_data["multi_criteria_optimizer"]["id"]
+ mco_data = json_data.get("multi_criteria_optimizer")
+ if mco_data is None:
+ # The file was saved without setting an MCO.
+ # The file is valid, we simply can't run any optimization yet.
+ return None
+
+ mco_id = mco_data["id"]
mco_bundle = registry.mco_bundle_by_id(mco_id)
return mco_bundle.create_model(
json_data["multi_criteria_optimizer"]["model_data"])
diff --git a/force_bdss/io/workflow_writer.py b/force_bdss/io/workflow_writer.py
index 05eff31..d0b2c8d 100644
--- a/force_bdss/io/workflow_writer.py
+++ b/force_bdss/io/workflow_writer.py
@@ -5,31 +5,33 @@ from ..bundle_registry_plugin import BundleRegistryPlugin
class WorkflowWriter(HasStrictTraits):
- bundle_registry = Instance(BundleRegistryPlugin)
-
- def __init__(self, bundle_registry, *args, **kwargs):
- self.bundle_registry = bundle_registry
-
- super(WorkflowWriter, self).__init__(*args, **kwargs)
-
def write(self, workflow, f):
data = {
"version": "1",
+ "workflow": {}
}
- data["multi_criteria_optimizer"] = \
- workflow.multi_criteria_optimizer.__getstate__()
-
+ wf_data = data["workflow"]
+ wf_data["multi_criteria_optimizer"] = {
+ "id": workflow.multi_criteria_optimizer.bundle.id,
+ "model_data": workflow.multi_criteria_optimizer.__getstate__()
+ }
kpic_data = []
for kpic in workflow.kpi_calculators:
- kpic_data.append(kpic.__getstate__())
+ kpic_data.append({
+ "id": kpic.bundle.id,
+ "model_data": kpic.__getstate__()}
+ )
- data["kpi_calculators"] = kpic_data
+ wf_data["kpi_calculators"] = kpic_data
ds_data = []
for ds in workflow.data_sources:
- ds_data.append(ds.__getstate__())
+ ds_data.append({
+ "id": ds.bundle.id,
+ "model_data": ds.__getstate__()
+ })
- data["data_sources"] = ds_data
+ wf_data["data_sources"] = ds_data
json.dump(data, f)
diff --git a/force_bdss/workspecs/workflow.py b/force_bdss/workspecs/workflow.py
index 8fc6be7..2796f23 100644
--- a/force_bdss/workspecs/workflow.py
+++ b/force_bdss/workspecs/workflow.py
@@ -6,6 +6,6 @@ from ..mco.base_mco_model import BaseMCOModel
class Workflow(HasStrictTraits):
- multi_criteria_optimizer = Instance(BaseMCOModel)
+ multi_criteria_optimizer = Instance(BaseMCOModel, allow_none=True)
data_sources = List(BaseDataSourceModel)
kpi_calculators = List(BaseKPICalculatorModel)
--
GitLab