diff --git a/tf-psa-crypto/tests/suites/test_suite_bignum.function b/tf-psa-crypto/tests/suites/test_suite_bignum.function
index 3ac4e10ea6..5e18441d7e 100644
--- a/tf-psa-crypto/tests/suites/test_suite_bignum.function
+++ b/tf-psa-crypto/tests/suites/test_suite_bignum.function
@@ -989,7 +989,13 @@ void mpi_exp_mod_min_RR(char *input_A, char *input_E,
      * against a smaller RR. */
     TEST_LE_U(RR.n, N.n - 1);
 
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    mbedtls_mpi_optionally_safe_codepath_reset();
+#endif
     res = mbedtls_mpi_exp_mod(&Z, &A, &E, &N, &RR);
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+#endif
     /* We know that exp_mod internally needs RR to be as large as N.
      * Validate that it is the case now, otherwise there was probably
      * a buffer overread. */
@@ -1022,7 +1028,26 @@ void mpi_exp_mod(char *input_A, char *input_E,
     TEST_ASSERT(mbedtls_test_read_mpi(&N, input_N) == 0);
     TEST_ASSERT(mbedtls_test_read_mpi(&X, input_X) == 0);
 
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    mbedtls_mpi_optionally_safe_codepath_reset();
+#endif
     res = mbedtls_mpi_exp_mod(&Z, &A, &E, &N, NULL);
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+#endif
+    TEST_ASSERT(res == exp_result);
+    if (res == 0) {
+        TEST_ASSERT(sign_is_valid(&Z));
+        TEST_ASSERT(mbedtls_mpi_cmp_mpi(&Z, &X) == 0);
+    }
+
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    mbedtls_mpi_optionally_safe_codepath_reset();
+#endif
+    res = mbedtls_mpi_exp_mod_unsafe(&Z, &A, &E, &N, NULL);
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_PUBLIC);
+#endif
     TEST_ASSERT(res == exp_result);
     if (res == 0) {
         TEST_ASSERT(sign_is_valid(&Z));
@@ -1030,7 +1055,13 @@ void mpi_exp_mod(char *input_A, char *input_E,
     }
 
     /* Now test again with the speed-up parameter supplied as an output. */
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    mbedtls_mpi_optionally_safe_codepath_reset();
+#endif
     res = mbedtls_mpi_exp_mod(&Z, &A, &E, &N, &RR);
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+#endif
     TEST_ASSERT(res == exp_result);
     if (res == 0) {
         TEST_ASSERT(sign_is_valid(&Z));
@@ -1038,7 +1069,13 @@ void mpi_exp_mod(char *input_A, char *input_E,
     }
 
     /* Now test again with the speed-up parameter supplied in calculated form. */
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    mbedtls_mpi_optionally_safe_codepath_reset();
+#endif
     res = mbedtls_mpi_exp_mod(&Z, &A, &E, &N, &RR);
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+#endif
     TEST_ASSERT(res == exp_result);
     if (res == 0) {
         TEST_ASSERT(sign_is_valid(&Z));
@@ -1078,7 +1115,21 @@ void mpi_exp_mod_size(int A_bytes, int E_bytes, int N_bytes,
         TEST_ASSERT(mbedtls_test_read_mpi(&RR, input_RR) == 0);
     }
 
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    mbedtls_mpi_optionally_safe_codepath_reset();
+#endif
     TEST_ASSERT(mbedtls_mpi_exp_mod(&Z, &A, &E, &N, &RR) == exp_result);
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+#endif
+
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    mbedtls_mpi_optionally_safe_codepath_reset();
+#endif
+    TEST_ASSERT(mbedtls_mpi_exp_mod_unsafe(&Z, &A, &E, &N, &RR) == exp_result);
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+#endif
 
 exit:
     mbedtls_mpi_free(&A); mbedtls_mpi_free(&E); mbedtls_mpi_free(&N);