diff --git a/tests/test_draw.py b/tests/test_draw.py
index 07cbc8aad3..d3cf7b5053 100644
--- a/tests/test_draw.py
+++ b/tests/test_draw.py
@@ -1259,7 +1259,7 @@ def test_fft():
                              [4,6,8,4],
                              [2,4,6,6] ],
                            xmin=-2, ymin=-2, dtype=dt, scale=0.1)
-        if is_gjax_galsim() and dt not in [np.complex128, complex]:
+        if is_jax_galsim() and dt not in [np.complex128, complex]:
             kim = xim.calculate_fft()
             xim2 = kim.calculate_inverse_fft()
             np.testing.assert_array_almost_equal(xim.array, xim2.array)
@@ -1298,7 +1298,7 @@ def test_fft():
         xim2 = galsim.Image([ [2,4,6],
                               [4,6,8] ],
                             xmin=-2, ymin=-1, dtype=dt, scale=0.1)
-        if is_gjax_galsim() and dt not in [np.complex128, complex]:
+        if is_jax_galsim() and dt not in [np.complex128, complex]:
             kim = xim.calculate_fft()
             kim2 = xim2.calculate_fft()
             np.testing.assert_array_almost_equal(kim.array, kim2.array)