diff --git a/tests/suites/test_suite_ecjpake.function b/tests/suites/test_suite_ecjpake.function
index ab6737f..449b368 100644
--- a/tests/suites/test_suite_ecjpake.function
+++ b/tests/suites/test_suite_ecjpake.function
@@ -1,7 +1,8 @@
 /* BEGIN_HEADER */
 #include "mbedtls/ecjpake.h"
+#include "legacy_or_psa.h"
 
-#if defined(MBEDTLS_ECP_DP_SECP256R1_ENABLED) && defined(MBEDTLS_SHA256_C)
+#if defined(MBEDTLS_ECP_DP_SECP256R1_ENABLED) && defined(MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA)
 static const unsigned char ecjpake_test_x1[] = {
     0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
     0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18,
@@ -90,7 +91,7 @@
 }
 
 #define ADD_SIZE( x )   x, sizeof( x )
-#endif /* MBEDTLS_ECP_DP_SECP256R1_ENABLED && MBEDTLS_SHA256_C */
+#endif /* MBEDTLS_ECP_DP_SECP256R1_ENABLED && MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA */
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -126,7 +127,7 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_SHA256_C */
+/* BEGIN_CASE depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA */
 void read_bad_md( data_t *msg )
 {
     mbedtls_ecjpake_context corrupt_ctx;
@@ -147,7 +148,7 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_SHA256_C */
+/* BEGIN_CASE depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA */
 void read_round_one( int role, data_t * msg, int ref_ret )
 {
     mbedtls_ecjpake_context ctx;
@@ -166,7 +167,7 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_SHA256_C */
+/* BEGIN_CASE depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA */
 void read_round_two_cli( data_t * msg, int ref_ret )
 {
     mbedtls_ecjpake_context ctx;
@@ -191,7 +192,7 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_SHA256_C */
+/* BEGIN_CASE depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_HAS_ALG_SHA_256_VIA_MD_OR_PSA */
 void read_round_two_srv( data_t * msg, int ref_ret )
 {
     mbedtls_ecjpake_context ctx;