3 files changed, 7 insertions, 5 deletions
diff --git a/test/test_parameters.py b/test/test_parameters.py
index 22efccf..e36b1a1 100755
--- a/test/test_parameters.py
+++ b/test/test_parameters.py
@@ -1,9 +1,9 @@
 #!/usr/bin/env python3
 
-from dfatool import dfatool as dt
 from dfatool import parameters
 from dfatool.utils import by_name_to_by_param
 from dfatool.functions import analytic
+from dfatool.model import ParallelParamFit
 import unittest
 
 import numpy as np
@@ -52,7 +52,7 @@ class TestModels(unittest.TestCase):
 
         # Fit individual functions for each parameter (only "p_linear" in this case)
 
-        paramfit = dt.ParallelParamFit(by_param)
+        paramfit = ParallelParamFit(by_param)
         paramfit.enqueue("TX", "power", 1, "p_linear")
         paramfit.fit()
 
@@ -133,7 +133,7 @@ class TestModels(unittest.TestCase):
         self.assertEqual(stats.depends_on_param("someKey", "ll", "log_inv"), True)
         self.assertEqual(stats.depends_on_param("someKey", "ll", "square_none"), False)
 
-        paramfit = dt.ParallelParamFit(by_param)
+        paramfit = ParallelParamFit(by_param)
         paramfit.enqueue("someKey", "lls", 0, "lin_lin")
         paramfit.enqueue("someKey", "lls", 1, "log_inv")
         paramfit.enqueue("someKey", "lls", 2, "square_none")
diff --git a/test/test_ptamodel.py b/test/test_ptamodel.py
index 9abe3c0..e153280 100755
--- a/test/test_ptamodel.py
+++ b/test/test_ptamodel.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
-from dfatool.dfatool import PTAModel, RawData, pta_trace_to_aggregate
+from dfatool.dfatool import RawData, pta_trace_to_aggregate
+from dfatool.model import PTAModel
 import os
 import unittest
 import pytest
diff --git a/test/test_timingharness.py b/test/test_timingharness.py
index 61cdf4b..cc4b766 100755
--- a/test/test_timingharness.py
+++ b/test/test_timingharness.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
-from dfatool.dfatool import AnalyticModel, TimingData, pta_trace_to_aggregate
+from dfatool.dfatool import TimingData, pta_trace_to_aggregate
+from dfatool.model import AnalyticModel
 from dfatool.parameters import prune_dependent_parameters
 import unittest