diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index e58abf9..361308b 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -1955,6 +1955,14 @@
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA1_C
 derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_1):"0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c":"":"":255 * 20:"2c91117204d745f3500d636a62f64f0ab3bae548aa53d423b0d1f27ebba6f5e5673a081d70cce7acfc48":""
 
+PSA key derivation: HKDF SHA-256, request too much capacity
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_set_capacity:PSA_ALG_HKDF(PSA_ALG_SHA_256):255 * 32 + 1:PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: HKDF SHA-1, request too much capacity
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA1_C
+derive_set_capacity:PSA_ALG_HKDF(PSA_ALG_SHA_1):255 * 20 + 1:PSA_ERROR_INVALID_ARGUMENT
+
 PSA key derivation: TLS 1.2 PSK-to-MS, SHA-256, PSK too long (160 Bytes)
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C:!PSA_PRE_1_0_KEY_DERIVATION
 derive_input:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):PSA_KEY_TYPE_DERIVE:PSA_KEY_DERIVATION_INPUT_SEED:"":PSA_KEY_DERIVATION_INPUT_SECRET:"01020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708":PSA_KEY_DERIVATION_INPUT_LABEL:"":PSA_SUCCESS:PSA_ERROR_INVALID_ARGUMENT:PSA_ERROR_BAD_STATE
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 316f2ed..858356d 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4053,6 +4053,28 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
+void derive_set_capacity( int alg_arg, int capacity_arg,
+                          int expected_status_arg )
+{
+    psa_algorithm_t alg = alg_arg;
+    size_t capacity = capacity_arg;
+    psa_status_t expected_status = expected_status_arg;
+    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
+
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
+
+    TEST_EQUAL( psa_key_derivation_set_capacity( &operation, capacity ),
+                expected_status );
+
+exit:
+    psa_key_derivation_abort( &operation );
+    PSA_DONE( );
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
 void derive_input( int alg_arg,
                    int key_type_arg,
                    int step1_arg, data_t *input1,