diff --git a/autolens/analysis/analysis/dataset.py b/autolens/analysis/analysis/dataset.py
index 76224bb55..82d107caf 100644
--- a/autolens/analysis/analysis/dataset.py
+++ b/autolens/analysis/analysis/dataset.py
@@ -31,6 +31,7 @@
 from autolens.analysis.positions import PositionsLH
 
 from autolens import exc
+from autofit.non_linear.test_mode import is_test_mode
 
 logger = logging.getLogger(__name__)
 
@@ -104,7 +105,7 @@ def __init__(
             raise_inversion_positions_likelihood_exception
         )
 
-        if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
+        if is_test_mode():
             self.raise_inversion_positions_likelihood_exception = False
 
         # Can be deleted after relevent AutoFIT PR merged
diff --git a/autolens/analysis/result.py b/autolens/analysis/result.py
index 133ff62d3..f147d2cb3 100644
--- a/autolens/analysis/result.py
+++ b/autolens/analysis/result.py
@@ -31,6 +31,7 @@
 )
 from autolens.lens.tracer import Tracer
 from autolens.point.solver import PointSolver
+from autofit.non_linear.test_mode import is_test_mode
 
 logger = logging.getLogger(__name__)
 
@@ -303,7 +304,7 @@ def positions_likelihood_from(
         The `PositionsLH` object used to apply a likelihood penalty or resample the positions.
         """
 
-        if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
+        if is_test_mode():
             return
 
         positions = (
diff --git a/autolens/quantity/model/visualizer.py b/autolens/quantity/model/visualizer.py
index b31454271..f29fd474d 100644
--- a/autolens/quantity/model/visualizer.py
+++ b/autolens/quantity/model/visualizer.py
@@ -1,6 +1,7 @@
 import os
 
 import autofit as af
+from autofit.non_linear.test_mode import is_test_mode
 
 from autogalaxy.quantity.model.plotter import PlotterQuantity
 
@@ -39,7 +40,7 @@ def visualize(
             via a non-linear search).
         """
 
-        if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
+        if is_test_mode():
             return
 
         fit = analysis.fit_quantity_for_instance(instance=instance)