diff --git a/force_bdss/tests/probe_classes/data_source.py b/force_bdss/tests/probe_classes/data_source.py
index 8cc11cb7117f4cdf00adab313b8e4dbf4ea5f8a3..18fcceecd758dedfa2563fdd07aed229b39cb068 100644
--- a/force_bdss/tests/probe_classes/data_source.py
+++ b/force_bdss/tests/probe_classes/data_source.py
@@ -1,4 +1,4 @@
-from traits.api import Bool, Function, Str, Int, on_trait_change
+from traits.api import Bool, Function, Str, Int, on_trait_change, Type
 
 from force_bdss.api import (
     BaseDataSourceFactory, BaseDataSourceModel, BaseDataSource,
@@ -47,7 +47,7 @@ class ProbeDataSourceFactory(BaseDataSourceFactory,
     id = Str('enthought.test.data_source')
     name = Str('test_data_source')
 
-    model_class = ProbeDataSourceModel
+    model_class = Type(ProbeDataSourceModel)
 
     def create_model(self, model_data=None):
         return self.model_class(
diff --git a/force_bdss/tests/probe_classes/kpi_calculator.py b/force_bdss/tests/probe_classes/kpi_calculator.py
index c3e455597371ca8f8b614cb44fcab5f9eac991aa..951ce9f252280297d4cc03ba9112d2fba602377c 100644
--- a/force_bdss/tests/probe_classes/kpi_calculator.py
+++ b/force_bdss/tests/probe_classes/kpi_calculator.py
@@ -1,4 +1,4 @@
-from traits.api import Bool, Function, Str, Int, on_trait_change
+from traits.api import Bool, Function, Str, Int, on_trait_change, Type
 
 from force_bdss.api import (
     BaseKPICalculatorFactory, BaseKPICalculatorModel, BaseKPICalculator,
@@ -47,7 +47,7 @@ class ProbeKPICalculatorFactory(BaseKPICalculatorFactory,
     id = Str('enthought.test.kpi_calculator')
     name = Str('test_kpi_calculator')
 
-    model_class = ProbeKPICalculatorModel
+    model_class = Type(ProbeKPICalculatorModel)
 
     def create_model(self, model_data=None):
         return self.model_class(