diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index 7f93496..48029b4 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -506,7 +506,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));
 
@@ -548,7 +548,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));
@@ -838,7 +838,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))) {
@@ -881,7 +881,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),
@@ -894,7 +894,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 761d877..f70b89a 100644
--- a/tests/src/test_helpers/ssl_helpers.c
+++ b/tests/src/test_helpers/ssl_helpers.c
@@ -91,7 +91,7 @@
     opts->resize_buffers = 1;
 #if defined(MBEDTLS_SSL_CACHE_C)
     opts->cache = NULL;
-    ASSERT_ALLOC(opts->cache, 1);
+    TEST_CALLOC_OR_FAIL(opts->cache, 1);
     mbedtls_ssl_cache_init(opts->cache);
 #if defined(MBEDTLS_HAVE_TIME)
     TEST_EQUAL(mbedtls_ssl_cache_get_timeout(opts->cache),
@@ -627,9 +627,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);