diff --git a/tf-psa-crypto/tests/suites/test_suite_bignum_core.function b/tf-psa-crypto/tests/suites/test_suite_bignum_core.function
index d5cc08e..c755287 100644
--- a/tf-psa-crypto/tests/suites/test_suite_bignum_core.function
+++ b/tf-psa-crypto/tests/suites/test_suite_bignum_core.function
@@ -1303,56 +1303,46 @@
     TEST_CF_SECRET(A, A_limbs * sizeof(mbedtls_mpi_uint));
     TEST_CF_SECRET(N, N_limbs * sizeof(mbedtls_mpi_uint));
     TEST_CF_SECRET(E, E_limbs * sizeof(mbedtls_mpi_uint));
-#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_codepath_reset();
-#endif
+
     mbedtls_mpi_core_exp_mod(Y, A, N, N_limbs, E, E_limbs, R2, T);
-#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
-#endif
-    TEST_EQUAL(0, memcmp(X, Y, N_limbs * sizeof(mbedtls_mpi_uint)));
 
     TEST_CF_PUBLIC(Y, N_limbs * sizeof(mbedtls_mpi_uint));
-
     TEST_EQUAL(0, memcmp(X, Y, N_limbs * sizeof(mbedtls_mpi_uint)));
 
     /* Test the unsafe variant */
+    TEST_CF_PUBLIC(A, A_limbs * sizeof(mbedtls_mpi_uint));
+    TEST_CF_PUBLIC(N, N_limbs * sizeof(mbedtls_mpi_uint));
+    TEST_CF_PUBLIC(E, E_limbs * sizeof(mbedtls_mpi_uint));
 
-#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_codepath_reset();
-#endif
     mbedtls_mpi_core_exp_mod_unsafe(Y, A, N, N_limbs, E, E_limbs, R2, T);
-#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_PUBLIC);
-#endif
+
     TEST_EQUAL(0, memcmp(X, Y, N_limbs * sizeof(mbedtls_mpi_uint)));
 
-    /* Check both with output aliased to input */
+    /*
+     * Check both with output aliased to input
+     */
 
     TEST_CALLOC(A_copy, A_limbs);
-    memcpy(A_copy, A, sizeof(*A_copy) * A_limbs);
+    memcpy(A_copy, A, sizeof(*A_copy) * A_limbs); // save A
 
+    /* Safe */
     TEST_CF_SECRET(A, A_limbs * sizeof(mbedtls_mpi_uint));
     TEST_CF_SECRET(N, N_limbs * sizeof(mbedtls_mpi_uint));
     TEST_CF_SECRET(E, E_limbs * sizeof(mbedtls_mpi_uint));
-#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_codepath_reset();
-#endif
+
     mbedtls_mpi_core_exp_mod(A, A, N, N_limbs, E, E_limbs, R2, T);
-#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
-#endif
-    TEST_EQUAL(0, memcmp(X, A, N_limbs * sizeof(mbedtls_mpi_uint)));
 
     TEST_CF_PUBLIC(A, A_limbs * sizeof(mbedtls_mpi_uint));
-    memcpy(A, A_copy, sizeof(*A) * A_limbs);
-#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_codepath_reset();
-#endif
+    TEST_EQUAL(0, memcmp(X, A, N_limbs * sizeof(mbedtls_mpi_uint)));
+
+    /* Unsafe */
+    memcpy(A, A_copy, sizeof(*A) * A_limbs); // restore A
+    TEST_CF_PUBLIC(A, A_limbs * sizeof(mbedtls_mpi_uint));
+    TEST_CF_PUBLIC(N, N_limbs * sizeof(mbedtls_mpi_uint));
+    TEST_CF_PUBLIC(E, E_limbs * sizeof(mbedtls_mpi_uint));
+
     mbedtls_mpi_core_exp_mod_unsafe(A, A, N, N_limbs, E, E_limbs, R2, T);
-#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_PUBLIC);
-#endif
+
     TEST_EQUAL(0, memcmp(X, A, N_limbs * sizeof(mbedtls_mpi_uint)));
 
 exit: