diff --git a/tests/include/test/psa_exercise_key.h b/tests/include/test/psa_exercise_key.h
index 0f3ee3d..bd523d1 100644
--- a/tests/include/test/psa_exercise_key.h
+++ b/tests/include/test/psa_exercise_key.h
@@ -46,11 +46,11 @@
  *
  * For simplicity's sake, stick to block ciphers with 16-byte blocks.
  */
-#if defined(MBEDTLS_AES_C)
+#if defined(PSA_WANT_KEY_TYPE_AES)
 #define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_AES
-#elif defined(MBEDTLS_ARIA_C)
+#elif defined(PSA_WANT_KEY_TYPE_ARIA)
 #define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_ARIA
-#elif defined(MBEDTLS_CAMELLIA_C)
+#elif defined(PSA_WANT_KEY_TYPE_CAMELLIA)
 #define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_CAMELLIA
 #undef KNOWN_SUPPORTED_BLOCK_CIPHER
 #endif
@@ -81,13 +81,13 @@
  *
  * This is used in some smoke tests.
  */
-#if defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CTR)
+#if defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(PSA_WANT_ALG_CTR)
 #define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CTR
-#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CBC)
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(PSA_WANT_ALG_CBC_NO_PADDING)
 #define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CBC_NO_PADDING
-#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CFB)
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(PSA_WANT_ALG_CFB)
 #define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CFB
-#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_OFB)
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(PSA_WANT_ALG_OFB)
 #define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_OFB
 #else
 #undef KNOWN_SUPPORTED_BLOCK_CIPHER_ALG