diff --git a/force_bdss/data_sources/base_data_source_model.py b/force_bdss/data_sources/base_data_source_model.py index 75f7c5396c13e869f270bd22bc5713080dbfb624..2326075deff38378a0aa868cb465cb7a2ef4db21 100644 --- a/force_bdss/data_sources/base_data_source_model.py +++ b/force_bdss/data_sources/base_data_source_model.py @@ -34,3 +34,4 @@ class BaseDataSourceModel(ABCHasStrictTraits): state["input_slot_maps"] = [ x.__getstate__() for x in self.input_slot_maps ] + return state diff --git a/force_bdss/data_sources/tests/test_base_data_source_model.py b/force_bdss/data_sources/tests/test_base_data_source_model.py new file mode 100644 index 0000000000000000000000000000000000000000..e0e3a0f9e63115f062561f2490ba12a0c01aff82 --- /dev/null +++ b/force_bdss/data_sources/tests/test_base_data_source_model.py @@ -0,0 +1,57 @@ +import unittest + +from force_bdss.core.input_slot_map import InputSlotMap + +try: + import mock +except ImportError: + from unittest import mock + +from force_bdss.data_sources.base_data_source_bundle import \ + BaseDataSourceBundle +from force_bdss.data_sources.base_data_source_model import BaseDataSourceModel + + +class DummyDataSourceModel(BaseDataSourceModel): + pass + + +class TestBaseDataSourceModel(unittest.TestCase): + def test_getstate(self): + model = DummyDataSourceModel(mock.Mock(spec=BaseDataSourceBundle)) + self.assertEqual( + model.__getstate__(), + { + "__traits_version__": "4.6.0", + "input_slot_maps": [], + "output_slot_names": [] + }) + + model.input_slot_maps = [ + InputSlotMap( + name="foo" + ), + InputSlotMap( + name="bar" + ) + ] + model.output_slot_names = ["baz", "quux"] + + self.assertEqual( + model.__getstate__(), + { + "__traits_version__": "4.6.0", + "input_slot_maps": [ + { + "__traits_version__": "4.6.0", + "source": "Environment", + "name": "foo" + }, + { + "__traits_version__": "4.6.0", + "source": "Environment", + "name": "bar" + } + ], + "output_slot_names": ["baz", "quux"] + })