diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index f8ed784..ff43780 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -505,7 +505,7 @@
     key_bits = psa_get_key_bits(&attributes);
     public_key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(private_key_type);
     public_key_length = PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE(public_key_type, key_bits);
-    ASSERT_ALLOC(public_key, public_key_length);
+    TEST_CALLOC_OR_FAIL(public_key, public_key_length);
     PSA_ASSERT(psa_export_public_key(key, public_key, public_key_length,
                                      &public_key_length));
 
@@ -547,7 +547,7 @@
     key_bits = psa_get_key_bits(&attributes);
     public_key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(private_key_type);
     public_key_length = PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE(public_key_type, key_bits);
-    ASSERT_ALLOC(public_key, public_key_length);
+    TEST_CALLOC_OR_FAIL(public_key, public_key_length);
     PSA_ASSERT(psa_export_public_key(key,
                                      public_key, public_key_length,
                                      &public_key_length));
@@ -807,7 +807,7 @@
     exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(
         psa_get_key_type(&attributes),
         psa_get_key_bits(&attributes));
-    ASSERT_ALLOC(exported, exported_size);
+    TEST_CALLOC_OR_FAIL(exported, exported_size);
 
     if ((usage & PSA_KEY_USAGE_EXPORT) == 0 &&
         !PSA_KEY_TYPE_IS_PUBLIC_KEY(psa_get_key_type(&attributes))) {
@@ -850,7 +850,7 @@
         exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(
             psa_get_key_type(&attributes),
             psa_get_key_bits(&attributes));
-        ASSERT_ALLOC(exported, exported_size);
+        TEST_CALLOC_OR_FAIL(exported, exported_size);
 
         TEST_EQUAL(psa_export_public_key(key, exported,
                                          exported_size, &exported_length),
@@ -863,7 +863,7 @@
         psa_get_key_type(&attributes));
     exported_size = PSA_EXPORT_PUBLIC_KEY_OUTPUT_SIZE(public_type,
                                                       psa_get_key_bits(&attributes));
-    ASSERT_ALLOC(exported, exported_size);
+    TEST_CALLOC_OR_FAIL(exported, exported_size);
 
     PSA_ASSERT(psa_export_public_key(key,
                                      exported, exported_size,
diff --git a/tests/src/test_helpers/ssl_helpers.c b/tests/src/test_helpers/ssl_helpers.c
index bd1f46c..e346ca0 100644
--- a/tests/src/test_helpers/ssl_helpers.c
+++ b/tests/src/test_helpers/ssl_helpers.c
@@ -566,9 +566,9 @@
     }
 
     cert = &(ep->cert);
-    ASSERT_ALLOC(cert->ca_cert, 1);
-    ASSERT_ALLOC(cert->cert, 1);
-    ASSERT_ALLOC(cert->pkey, 1);
+    TEST_CALLOC_OR_FAIL(cert->ca_cert, 1);
+    TEST_CALLOC_OR_FAIL(cert->cert, 1);
+    TEST_CALLOC_OR_FAIL(cert->pkey, 1);
 
     mbedtls_x509_crt_init(cert->ca_cert);
     mbedtls_x509_crt_init(cert->cert);