diff --git a/tests/suites/test_suite_hkdf.function b/tests/suites/test_suite_hkdf.function
index feb1717..1e112ba 100644
--- a/tests/suites/test_suite_hkdf.function
+++ b/tests/suites/test_suite_hkdf.function
@@ -44,7 +44,7 @@
     TEST_ASSERT( md != NULL );
 
     output_prk_len = mbedtls_md_get_size( md );
-    output_prk = mbedtls_calloc( 1, output_prk_len );
+    ASSERT_ALLOC( output_prk, output_prk_len );
 
     ikm = mbedtls_test_unhexify_alloc( hex_ikm_string, &ikm_len );
     salt = mbedtls_test_unhexify_alloc( hex_salt_string, &salt_len );
@@ -113,7 +113,7 @@
     fake_md_info.type = MBEDTLS_MD_NONE;
     fake_md_info.size = hash_len;
 
-    prk = mbedtls_calloc( MBEDTLS_MD_MAX_SIZE, 1 );
+    ASSERT_ALLOC( prk, MBEDTLS_MD_MAX_SIZE);
     salt_len = 0;
     ikm_len = 0;
 
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index c8b70a3..7149840 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -3817,7 +3817,7 @@
     size_t ikm_len, salt_len, prk_len, output_prk_size, output_prk_len;
 
     output_prk_size = PSA_HASH_LENGTH( alg );
-    output_prk = mbedtls_calloc( 1, output_prk_size );
+    ASSERT_ALLOC( output_prk, output_prk_size );
 
     ikm = mbedtls_test_unhexify_alloc( hex_ikm_string, &ikm_len );
     salt = mbedtls_test_unhexify_alloc( hex_salt_string, &salt_len );
@@ -3850,8 +3850,8 @@
     unsigned char *prk = NULL;
     size_t salt_len, ikm_len, prk_len;
 
-    prk = mbedtls_calloc( PSA_MAC_MAX_SIZE, 1 );
     salt_len = hash_len;
+    ASSERT_ALLOC( prk, PSA_MAC_MAX_SIZE);
     ikm_len = 0;
     prk_len = 0;