diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 718b10c..eff07ad 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -3454,6 +3454,10 @@
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:PSA_WANT_ECC_SECP_R1_256
 derive_input:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":PSA_SUCCESS:0:UNUSED:"":UNUSED:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
 
+PSA key derivation: reject calling input functions without calling setup
+depends_on:PSA_WANT_ALG_SHA_256
+derive_input:0:PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_DERIVE:"61206c6162656c":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
+
 PSA key derivation over capacity: HKDF
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_over_capacity:PSA_ALG_HKDF(PSA_ALG_SHA_256)
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 838717e..39dd89a 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4520,7 +4520,9 @@
     psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_DERIVE);
     psa_set_key_algorithm(&attributes, alg);
 
-    PSA_ASSERT(psa_key_derivation_setup(&operation, alg));
+    if (alg != 0) {
+        PSA_ASSERT(psa_key_derivation_setup(&operation, alg));
+    }
 
     for (i = 0; i < ARRAY_LENGTH(steps); i++) {
         mbedtls_test_set_step(i);
@@ -4566,16 +4568,6 @@
     }
     TEST_EQUAL(actual_output_status, expected_output_status);
 
-    /* Test calling input functions after operation has been aborted
-       result in PSA_ERROR_BAD_STATE error.
-    */
-    psa_key_derivation_abort(&operation);
-
-    TEST_EQUAL(psa_key_derivation_input_bytes(
-                    &operation, steps[0],
-                    inputs[0]->x, inputs[0]->len),
-                PSA_ERROR_BAD_STATE);
-
 exit:
     psa_key_derivation_abort(&operation);
     for (i = 0; i < ARRAY_LENGTH(keys); i++) {