diff --git a/tests/include/test/psa_exercise_key.h b/tests/include/test/psa_exercise_key.h
index fa57d88..713b093 100644
--- a/tests/include/test/psa_exercise_key.h
+++ b/tests/include/test/psa_exercise_key.h
@@ -123,6 +123,9 @@
  * \param input2            The first input to pass.
  * \param input2_length     The length of \p input2 in bytes.
  * \param capacity          The capacity to set.
+ * \param key_destroyable   If set to 1, a failure due to the key not existing
+ *                          or the key being destroyed mid-operation will only
+ *                          be reported if the error code is unexpected.
  *
  * \return                  \c 1 on success, \c 0 on failure.
  */
@@ -132,7 +135,7 @@
     psa_algorithm_t alg,
     const unsigned char *input1, size_t input1_length,
     const unsigned char *input2, size_t input2_length,
-    size_t capacity);
+    size_t capacity, int key_destroyable);
 
 /** Perform a key agreement using the given key pair against its public key
  * using psa_raw_key_agreement().
diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index 4700739..7260f1a 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -489,16 +489,22 @@
     psa_algorithm_t alg,
     const unsigned char *input1, size_t input1_length,
     const unsigned char *input2, size_t input2_length,
-    size_t capacity)
+    size_t capacity, int key_destroyable)
 {
     PSA_ASSERT(psa_key_derivation_setup(operation, alg));
+    psa_status_t status = PSA_SUCCESS;
     if (PSA_ALG_IS_HKDF(alg)) {
         PSA_ASSERT(psa_key_derivation_input_bytes(operation,
                                                   PSA_KEY_DERIVATION_INPUT_SALT,
                                                   input1, input1_length));
-        PSA_ASSERT(psa_key_derivation_input_key(operation,
-                                                PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                key));
+        status = psa_key_derivation_input_key(operation,
+                                              PSA_KEY_DERIVATION_INPUT_SECRET,
+                                              key);
+        if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+            /* The key has been destroyed. */
+            return 1;
+        }
+        PSA_ASSERT(status);
         PSA_ASSERT(psa_key_derivation_input_bytes(operation,
                                                   PSA_KEY_DERIVATION_INPUT_INFO,
                                                   input2,
@@ -507,13 +513,23 @@
         PSA_ASSERT(psa_key_derivation_input_bytes(operation,
                                                   PSA_KEY_DERIVATION_INPUT_SALT,
                                                   input1, input1_length));
-        PSA_ASSERT(psa_key_derivation_input_key(operation,
-                                                PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                key));
+        status = psa_key_derivation_input_key(operation,
+                                              PSA_KEY_DERIVATION_INPUT_SECRET,
+                                              key);
+        if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+            /* The key has been destroyed. */
+            return 1;
+        }
+        PSA_ASSERT(status);
     } else if (PSA_ALG_IS_HKDF_EXPAND(alg)) {
-        PSA_ASSERT(psa_key_derivation_input_key(operation,
-                                                PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                key));
+        status = psa_key_derivation_input_key(operation,
+                                              PSA_KEY_DERIVATION_INPUT_SECRET,
+                                              key);
+        if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+            /* The key has been destroyed. */
+            return 1;
+        }
+        PSA_ASSERT(status);
         PSA_ASSERT(psa_key_derivation_input_bytes(operation,
                                                   PSA_KEY_DERIVATION_INPUT_INFO,
                                                   input2,
@@ -523,9 +539,14 @@
         PSA_ASSERT(psa_key_derivation_input_bytes(operation,
                                                   PSA_KEY_DERIVATION_INPUT_SEED,
                                                   input1, input1_length));
-        PSA_ASSERT(psa_key_derivation_input_key(operation,
-                                                PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                key));
+        status = psa_key_derivation_input_key(operation,
+                                              PSA_KEY_DERIVATION_INPUT_SECRET,
+                                              key);
+        if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+            /* The key has been destroyed. */
+            return 1;
+        }
+        PSA_ASSERT(status);
         PSA_ASSERT(psa_key_derivation_input_bytes(operation,
                                                   PSA_KEY_DERIVATION_INPUT_LABEL,
                                                   input2, input2_length));
@@ -537,9 +558,14 @@
                                                   PSA_KEY_DERIVATION_INPUT_SALT,
                                                   input2,
                                                   input2_length));
-        PSA_ASSERT(psa_key_derivation_input_key(operation,
-                                                PSA_KEY_DERIVATION_INPUT_PASSWORD,
-                                                key));
+        status = psa_key_derivation_input_key(operation,
+                                              PSA_KEY_DERIVATION_INPUT_PASSWORD,
+                                              key);
+        if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+            /* The key has been destroyed. */
+            return 1;
+        }
+        PSA_ASSERT(status);
     } else if (alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS) {
         PSA_ASSERT(psa_key_derivation_input_bytes(operation,
                                                   PSA_KEY_DERIVATION_INPUT_SECRET,
@@ -561,7 +587,8 @@
 
 static int exercise_key_derivation_key(mbedtls_svc_key_id_t key,
                                        psa_key_usage_t usage,
-                                       psa_algorithm_t alg)
+                                       psa_algorithm_t alg,
+                                       int key_destroyable)
 {
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     unsigned char input1[] = "Input 1";
@@ -575,14 +602,20 @@
         if (!mbedtls_test_psa_setup_key_derivation_wrap(&operation, key, alg,
                                                         input1, input1_length,
                                                         input2, input2_length,
-                                                        capacity)) {
+                                                        capacity, key_destroyable)) {
             goto exit;
         }
 
-        PSA_ASSERT(psa_key_derivation_output_bytes(&operation,
-                                                   output,
-                                                   capacity));
-        PSA_ASSERT(psa_key_derivation_abort(&operation));
+        psa_status_t status = psa_key_derivation_output_bytes(&operation,
+                                                              output,
+                                                              capacity);
+        if (key_destroyable && status == PSA_ERROR_BAD_STATE) {
+            /* The key has been destroyed. */
+            PSA_ASSERT(psa_key_derivation_abort(&operation));
+        } else {
+            PSA_ASSERT(status);
+            PSA_ASSERT(psa_key_derivation_abort(&operation));
+        }
     }
 
     return 1;
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;
     }