diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 233dbe6ad3..483c3203ce 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -11,7 +11,10 @@
 
 #include "test/helpers.h"
 
-#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
+#if ((MBEDTLS_VERSION_MAJOR < 4) \
+    && defined(MBEDTLS_PSA_CRYPTO_C)) \
+    || (MBEDTLS_VERSION_MAJOR >= 4 \
+        && defined(MBEDTLS_PSA_CRYPTO_CLIENT))
 #include "test/psa_helpers.h"
 #include <psa/crypto.h>
 #endif
@@ -40,7 +43,8 @@
         mbedtls_psa_crypto_free();                                      \
     }                                                                   \
     while (0)
-#elif defined(MBEDTLS_PSA_CRYPTO_CLIENT) /* MBEDTLS_PSA_CRYPTO_CLIENT && !MBEDTLS_PSA_CRYPTO_C */
+#elif MBEDTLS_VERSION_MAJOR >= 4 && \
+    defined(MBEDTLS_PSA_CRYPTO_CLIENT)   /* MBEDTLS_PSA_CRYPTO_CLIENT && !MBEDTLS_PSA_CRYPTO_C */
 #define PSA_INIT() PSA_ASSERT(psa_crypto_init())
 #define PSA_DONE() mbedtls_psa_crypto_free();
 #else  /* MBEDTLS_PSA_CRYPTO_CLIENT && !MBEDTLS_PSA_CRYPTO_C */
@@ -48,7 +52,10 @@
 #define PSA_DONE() ((void) 0)
 #endif /* MBEDTLS_PSA_CRYPTO_C */
 
-#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
+#if ((MBEDTLS_VERSION_MAJOR < 4) \
+    && defined(MBEDTLS_PSA_CRYPTO_C)) \
+    || (MBEDTLS_VERSION_MAJOR >= 4 \
+        && defined(MBEDTLS_PSA_CRYPTO_CLIENT))
 
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
 
@@ -319,7 +326,7 @@ uint64_t mbedtls_test_parse_binary_string(data_t *bin_string);
     }                                                                      \
     while (0)
 
-#endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
+#endif /* MBEDTLS_PSA_CRYPTO_CLIENT || MBEDTLS_PSA_CRYPTO_C */
 
 /** \def USE_PSA_INIT
  *
diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index a496de0212..5c821e2f85 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -13,7 +13,8 @@
 
 #if ((MBEDTLS_VERSION_MAJOR < 4) \
     && defined(MBEDTLS_PSA_CRYPTO_C)) \
-    || defined(MBEDTLS_PSA_CRYPTO_CLIENT)
+    || (MBEDTLS_VERSION_MAJOR >= 4 \
+        && defined(MBEDTLS_PSA_CRYPTO_CLIENT))
 
 #include <mbedtls/asn1.h>
 #include <psa/crypto.h>