Fix psa_key_derivation_input_integer() not detecting bad state

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 94bf28b..9516867 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -8893,6 +8893,25 @@
     }
     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);
+    for (i = 0; i < ARRAY_LENGTH(steps); i++) {
+        if (key_types[i] == INPUT_INTEGER) {
+            TEST_EQUAL(psa_key_derivation_input_integer(
+                    &operation, steps[i],
+                    mbedtls_test_parse_binary_string(inputs[i])),
+                    PSA_ERROR_BAD_STATE);
+            break;
+        }
+    }
+
+    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++) {