Fix issues with the run function

force_bdss/tests/probe_classes/data_source.py

@@ -19,14 +19,9 @@ class ProbeDataSource(BaseDataSource):
     run_called = Bool(False)
     slots_called = Bool(False)
 
-    def __init__(self, factory, run_function=None, *args, **kwargs):
-        if run_function is None:
-            self.run_function = run_func
-        super(ProbeDataSource, self).__init__(self, factory, *args, **kwargs)
-
     def run(self, model, parameters):
         self.run_called = True
-        self.run_function(model, parameters)
+        return self.run_function(model, parameters)
 
     def slots(self, model):
         self.slots_called = True

force_bdss/tests/probe_classes/evaluator_factory.py

@@ -8,15 +8,18 @@ from envisage.plugin import Plugin
 from traits.api import HasStrictTraits, Function, Str, Int
 
 
+def run_func(*args, **kwargs):
+    return []
+
+
 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
+    run_function = Function(default_value=run_func)
 
     input_slots_type = Str('PRESSURE')
     output_slots_type = Str('PRESSURE')

force_bdss/tests/probe_classes/kpi_calculator.py

@@ -14,20 +14,14 @@ def run_func(*args, **kwargs):
 
 
 class ProbeKPICalculator(BaseKPICalculator):
-    run_function = Function()
+    run_function = Function(default_value=run_func)
 
     run_called = Bool(False)
     slots_called = Bool(False)
 
-    def __init__(self, factory, run_function=None, *args, **kwargs):
-        if run_function is None:
-            self.run_function = run_func
-        super(ProbeKPICalculator, self).__init__(
-            self, factory, *args, **kwargs)
-
     def run(self, model, parameters):
         self.run_called = True
-        self.run_function(model, parameters)
+        return self.run_function(model, parameters)
 
     def slots(self, model):
         self.slots_called = True
@@ -39,11 +33,6 @@ class ProbeKPICalculator(BaseKPICalculator):
                   for _ in range(model.output_slots_size))
         )
 
-    def _run_function_default(self):
-        def run_func(*args, **kwargs):
-            pass
-        return run_func
-
 
 class ProbeKPICalculatorModel(BaseKPICalculatorModel):
     input_slots_type = Str('PRESSURE')

force_bdss/tests/probe_classes/mco.py

@@ -18,22 +18,13 @@ def run_func(*args, **kwargs):
 
 
 class ProbeMCO(BaseMCO):
-    run_function = Function()
+    run_function = Function(default_value=run_func)
 
     run_called = Bool(False)
 
-    def __init__(self, factory, run_function=None, *args, **kwargs):
-        if run_function is None:
-            self.run_function = run_func
-        super(ProbeMCO, self).__init__(self, factory, *args, **kwargs)
-
     def run(self, model):
         self.run_called = True
-
-    def _run_function_default(self):
-        def run_func(*args, **kwargs):
-            pass
-        return run_func
+        return self.run_function(model)
 
 
 class ProbeParameter(BaseMCOParameter):