Add a test for psa_key_derivation_input
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index b21a8f1..7954d33 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4065,6 +4065,65 @@
 }
 /* END_CASE */
 
+/* BEGIN_CASE */
+void derive_input( int alg_arg,
+                   int key_type_arg,
+                   int step1_arg, data_t *input1,
+                   int step2_arg, data_t *input2,
+                   int step3_arg, data_t *input3,
+                   int expected_status_arg1,
+                   int expected_status_arg2,
+                   int expected_status_arg3 )
+{
+    psa_algorithm_t alg = alg_arg;
+    size_t key_type = key_type_arg;
+    psa_key_derivation_step_t steps[] = {step1_arg, step2_arg, step3_arg};
+    psa_status_t expected_statuses[] = {expected_status_arg1,
+                                        expected_status_arg2,
+                                        expected_status_arg3};
+    data_t *inputs[] = {input1, input2, input3};
+    psa_key_handle_t handles[] = {0, 0, 0};
+    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    size_t i;
+
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, key_type );
+
+    PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
+
+    for( i = 0; i < ARRAY_LENGTH( steps ); i++ )
+    {
+        switch( steps[i] )
+        {
+            case PSA_KEY_DERIVATION_INPUT_SECRET:
+                PSA_ASSERT( psa_import_key( &attributes,
+                                            inputs[i]->x, inputs[i]->len,
+                                            &handles[i] ) );
+                TEST_EQUAL( psa_key_derivation_input_key( &operation, steps[i],
+                                                          handles[i] ),
+                            expected_statuses[i] );
+                break;
+            default:
+                TEST_EQUAL( psa_key_derivation_input_bytes(
+                                    &operation, steps[i],
+                                    inputs[i]->x, inputs[i]->len ),
+                            expected_statuses[i] );
+                break;
+        }
+    }
+
+exit:
+    psa_key_derivation_abort( &operation );
+    for( i = 0; i < ARRAY_LENGTH( handles ); i++ )
+        psa_destroy_key( handles[i] );
+    PSA_DONE( );
+}
+/* END_CASE */
+
 /* BEGIN_CASE depends_on:PSA_PRE_1_0_KEY_DERIVATION */
 void test_derive_invalid_key_derivation_state( )
 {