diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index dfddbb9..7ef8618 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -8780,7 +8780,7 @@
     if (!mbedtls_test_psa_setup_key_derivation_wrap(&operation, key, alg,
                                                     input1, input1_length,
                                                     input2, input2_length,
-                                                    capacity)) {
+                                                    capacity, 0)) {
         goto exit;
     }
 
@@ -9099,7 +9099,7 @@
     if (!mbedtls_test_psa_setup_key_derivation_wrap(&operation, key, alg,
                                                     input1->x, input1->len,
                                                     input2->x, input2->len,
-                                                    requested_capacity)) {
+                                                    requested_capacity, 0)) {
         goto exit;
     }
 
@@ -9216,7 +9216,7 @@
     if (!mbedtls_test_psa_setup_key_derivation_wrap(&operation, base_key, alg,
                                                     input1->x, input1->len,
                                                     input2->x, input2->len,
-                                                    capacity)) {
+                                                    capacity, 0)) {
         goto exit;
     }
 
@@ -9286,7 +9286,7 @@
     if (!mbedtls_test_psa_setup_key_derivation_wrap(&operation, base_key, alg,
                                                     input1->x, input1->len,
                                                     input2->x, input2->len,
-                                                    capacity)) {
+                                                    capacity, 0)) {
         goto exit;
     }
 
@@ -9299,7 +9299,7 @@
     if (!mbedtls_test_psa_setup_key_derivation_wrap(&operation, base_key, alg,
                                                     input1->x, input1->len,
                                                     input2->x, input2->len,
-                                                    capacity)) {
+                                                    capacity, 0)) {
         goto exit;
     }
 
@@ -9370,7 +9370,7 @@
             &operation, base_key, alg,
             input1->x, input1->len,
             input2->x, input2->len,
-            PSA_KEY_DERIVATION_UNLIMITED_CAPACITY) == 0) {
+            PSA_KEY_DERIVATION_UNLIMITED_CAPACITY, 0) == 0) {
         goto exit;
     }
 
@@ -9435,7 +9435,7 @@
             &operation, base_key, alg,
             input1->x, input1->len,
             input2->x, input2->len,
-            PSA_KEY_DERIVATION_UNLIMITED_CAPACITY) == 0) {
+            PSA_KEY_DERIVATION_UNLIMITED_CAPACITY, 0) == 0) {
         goto exit;
     }
 
@@ -9499,7 +9499,7 @@
     if (!mbedtls_test_psa_setup_key_derivation_wrap(&operation, base_key, alg,
                                                     input1->x, input1->len,
                                                     input2->x, input2->len,
-                                                    SIZE_MAX)) {
+                                                    SIZE_MAX, 0)) {
         goto exit;
     }