Create probe data source classes

force_bdss/tests/probe_classes/data_source.py

+from traits.api import Bool, Function, Str, Int, on_trait_change
+
+from force_bdss.api import (
+    BaseDataSourceFactory, BaseDataSourceModel, BaseDataSource,
+    Slot
+)
+
+from .evaluator_factory import ProbeEvaluatorFactory
+
+
+class ProbeDataSource(BaseDataSource):
+    run_function = Function
+
+    run_called = Bool(False)
+    slots_called = Bool(False)
+
+    def run(self, model, parameters):
+        self.run_called = True
+        self.run_function(model, parameters)
+
+    def slots(self, model):
+        self.slots_called = True
+        return (
+            tuple(Slot(type=model.input_slots_type)
+                  for _ in range(model.input_slots_size))
+        ), (
+            tuple(Slot(type=model.output_slots_type)
+                  for _ in range(model.output_slots_size))
+        )
+
+
+class ProbeDataSourceModel(BaseDataSourceModel):
+    input_slots_type = Str('PRESSURE')
+    output_slots_type = Str('PRESSURE')
+
+    input_slots_size = Int(0)
+    output_slots_size = Int(0)
+
+    @on_trait_change('input_slots_type,output_slots_type,'
+                     'input_slots_size,output_slots_size')
+    def update_slots(self):
+        self.changes_slots = True
+
+
+class ProbeDataSourceFactory(BaseDataSourceFactory,
+                             ProbeEvaluatorFactory):
+    id = Str('enthought.test.data_source')
+    name = Str('test_data_source')
+
+    model_class = ProbeDataSourceModel
+
+    def create_model(self, model_data=None):
+        return self.model_class(
+            factory=self,
+            input_slots_type=self.input_slots_type,
+            output_slots_type=self.output_slots_type,
+            input_slots_size=self.input_slots_size,
+            output_slots_size=self.output_slots_size,
+            **model_data
+        )
+
+    def create_data_source(self):
+        return ProbeDataSource(
+            factory=self,
+            run_function=self.run_function,
+        )

force_bdss/tests/probe_classes/evaluator_factory.py

+try:
+    import mock
+except ImportError:
+    from unittest import mock
+
+from envisage.plugin import Plugin
+
+from traits.api import HasStrictTraits, Function, Str, Int
+
+
+class ProbeEvaluatorFactory(HasStrictTraits):
+    def __init__(self, plugin=None, *args, **kwargs):
+        if plugin is None:
+            plugin = mock.Mock(Plugin)
+
+        super(ProbeEvaluatorFactory, self).__init__(
+            plugin=plugin, *args, **kwargs)
+
+    run_function = Function
+
+    input_slots_type = Str('PRESSURE')
+    output_slots_type = Str('PRESSURE')
+
+    input_slots_size = Int(0)
+    output_slots_size = Int(0)

force_bdss/tests/probe_classes/factory_registry_plugin.py

@@ -3,6 +3,7 @@ from traits.api import List
 from force_bdss.factory_registry_plugin import FactoryRegistryPlugin
 
 from .kpi_calculator import ProbeKPICalculatorFactory
+from .data_source import ProbeDataSourceFactory
 
 
 class ProbeFactoryRegistryPlugin(FactoryRegistryPlugin):
@@ -12,4 +13,7 @@ class ProbeFactoryRegistryPlugin(FactoryRegistryPlugin):
     notification_listener_factories = List()
 
     def _kpi_calculator_factories_default(self):
-        return ProbeKPICalculatorFactory(self)
+        return [ProbeKPICalculatorFactory(self)]
+
+    def _data_source_factories_default(self):
+        return [ProbeDataSourceFactory(self)]

force_bdss/tests/probe_classes/kpi_calculator.py

-try:
-    import mock
-except ImportError:
-    from unittest import mock
-
-from envisage.plugin import Plugin
-
-from traits.api import (HasStrictTraits, Bool, Function, Str, Int,
-                        on_trait_change)
+from traits.api import Bool, Function, Str, Int, on_trait_change
 
 from force_bdss.api import (
     BaseKPICalculatorFactory, BaseKPICalculatorModel, BaseKPICalculator,
     Slot
 )
 
+from .evaluator_factory import ProbeEvaluatorFactory
+
 
 class ProbeEvaluator(BaseKPICalculator):
     run_function = Function
@@ -35,23 +29,6 @@ class ProbeEvaluator(BaseKPICalculator):
         )
 
 
-class ProbeEvaluatorFactory(HasStrictTraits):
-    def __init__(self, plugin=None, *args, **kwargs):
-        if plugin is None:
-            plugin = mock.Mock(Plugin)
-
-        super(ProbeEvaluatorFactory, self).__init__(
-            plugin=plugin, *args, **kwargs)
-
-    run_function = Function
-
-    input_slots_type = Str('PRESSURE')
-    output_slots_type = Str('PRESSURE')
-
-    input_slots_size = Int(0)
-    output_slots_size = Int(0)
-
-
 class ProbeKPICalculatorModel(BaseKPICalculatorModel):
     input_slots_type = Str('PRESSURE')
     output_slots_type = Str('PRESSURE')