diff --git a/force_bdss/io/tests/test_workflow_writer.py b/force_bdss/io/tests/test_workflow_writer.py index f7a0ccc9cbdc31606b8ab323eb0ac3a8de76c94e..fc197c306921d8ff7fe6de1ecd6340249398b7c6 100644 --- a/force_bdss/io/tests/test_workflow_writer.py +++ b/force_bdss/io/tests/test_workflow_writer.py @@ -3,6 +3,10 @@ import unittest from six import StringIO +from force_bdss.data_sources.base_data_source_factory import \ + BaseDataSourceFactory +from force_bdss.data_sources.base_data_source_model import BaseDataSourceModel +from force_bdss.data_sources.i_data_source_factory import IDataSourceFactory from force_bdss.factory_registry_plugin import FactoryRegistryPlugin from force_bdss.io.workflow_reader import WorkflowReader from force_bdss.mco.parameters.base_mco_parameter import BaseMCOParameter @@ -36,6 +40,13 @@ class TestWorkflowWriter(unittest.TestCase): self.mock_registry.mco_factory_by_id = mock.Mock( return_value=mock_mco_factory) + datasource_factory = BaseDataSourceFactory( + id=factory_id("enthought", "mock2"), plugin=None) + + self.mock_registry.data_source_factory_by_id = mock.Mock( + return_value=datasource_factory + ) + def test_write(self): wfwriter = WorkflowWriter() fp = StringIO() @@ -58,6 +69,9 @@ class TestWorkflowWriter(unittest.TestCase): wf_result = wfreader.read(fp) self.assertEqual(wf_result.mco.factory.id, wf.mco.factory.id) + self.assertEqual(len(wf_result.execution_layers), 2) + self.assertEqual(len(wf_result.execution_layers[0]), 2) + self.assertEqual(len(wf_result.execution_layers[1]), 1) def _create_mock_workflow(self): wf = Workflow() @@ -73,6 +87,18 @@ class TestWorkflowWriter(unittest.TestCase): ) ) ] + wf.execution_layers = [ + [ + BaseDataSourceModel(mock.Mock(spec=IDataSourceFactory, + id=factory_id("enthought", "mock2"))), + BaseDataSourceModel(mock.Mock(spec=IDataSourceFactory, + id=factory_id("enthought", "mock2"))), + ], + [ + BaseDataSourceModel(mock.Mock(spec=IDataSourceFactory, + id=factory_id("enthought", "mock2"))) + ] + ] return wf def test_write_and_read_empty_workflow(self):