Refactor

force_bdss/tests/probe_classes.py

@@ -5,7 +5,8 @@ except ImportError:
 
 from envisage.plugin import Plugin
 
-from traits.api import HasStrictTraits, Bool, Function, Str, Int, Dict
+from traits.api import (HasStrictTraits, Bool, Function, Str, Int, Dict,
+                        on_trait_change)
 
 from force_bdss.api import (
     BaseKPICalculatorFactory, BaseKPICalculatorModel, BaseKPICalculator,
@@ -13,15 +14,9 @@ from force_bdss.api import (
 )
 
 
-class ProbeEvaluator(HasStrictTraits):
+class ProbeEvaluator(BaseKPICalculator):
     run_function = Function
 
-    input_slots_type = Str('PRESSURE')
-    output_slots_type = Str('PRESSURE')
-
-    input_slots_size = Int(0)
-    output_slots_size = Int(0)
-
     run_called = Bool(False)
     slots_called = Bool(False)
 
@@ -32,11 +27,11 @@ class ProbeEvaluator(HasStrictTraits):
     def slots(self, model):
         self.slots_called = True
         return (
-            tuple(Slot(type=self.input_slots_type)
-                  for _ in range(self.input_slots_size))
+            tuple(Slot(type=model.input_slots_type)
+                  for _ in range(model.input_slots_size))
         ), (
-            tuple(Slot(type=self.output_slots_type)
-                  for _ in range(self.output_slots_size))
+            tuple(Slot(type=model.output_slots_type)
+                  for _ in range(model.output_slots_size))
         )
 
 
@@ -59,17 +54,17 @@ class ProbeEvaluatorFactory(HasStrictTraits):
     model_data = Dict()
 
 
-class ProbeKPICalculator(BaseKPICalculator, ProbeEvaluator):
-    pass
-
-
 class ProbeKPICalculatorModel(BaseKPICalculatorModel):
-    def __init__(self, factory, model_data, *args, **kwargs):
-        for key, value in model_data.items():
-            setattr(self, key, value)
+    input_slots_type = Str('PRESSURE')
+    output_slots_type = Str('PRESSURE')
+
+    input_slots_size = Int(0)
+    output_slots_size = Int(0)
 
-        super(ProbeKPICalculatorModel, self).__init__(
-            self, factory, *args, **kwargs)
+    @on_trait_change('input_slots_type,output_slots_type,'
+                     'input_slots_size,output_slots_size')
+    def update_slots(self):
+        self.changes_slots = True
 
 
 class ProbeKPICalculatorFactory(BaseKPICalculatorFactory,
@@ -77,15 +72,19 @@ class ProbeKPICalculatorFactory(BaseKPICalculatorFactory,
     id = Str('enthought.test.kpi_calculator')
     name = Str('test_kpi_calculator')
 
-    def create_model(self, model_data=None):
-        return ProbeKPICalculatorModel(self, self.model_data)
+    model_class = ProbeKPICalculatorModel
 
-    def create_kpi_calculator(self):
-        return ProbeKPICalculator(
+    def create_model(self, model_data=None):
+        return self.model_class(
             factory=self,
-            run_function=self.run_function,
             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
         )
+
+    def create_kpi_calculator(self):
+        return ProbeEvaluator(
+            factory=self,
+            run_function=self.run_function,
+        )