derive_input test function: Try output afterwards

After passing some inputs, try getting one byte of output, just to
check that this succeeds (for a valid sequence of inputs) or fails
with BAD_STATE (for an invalid sequence of inputs). Either output a
1-byte key or a 1-byte buffer depending on the test data.

The test data was expanded as follows:
* Output key type (or not a key): same as the SECRET input if success
  is expected, otherwise NONE.
* Expected status: PSA_SUCCESS after valid inputs, BAD_STATE after any
  invalid input.
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 11b17bc..87529ac 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4278,7 +4278,8 @@
                    int step_arg2, int key_type_arg2, data_t *input2,
                    int expected_status_arg2,
                    int step_arg3, int key_type_arg3, data_t *input3,
-                   int expected_status_arg3 )
+                   int expected_status_arg3,
+                   int output_key_type_arg, int expected_output_status_arg )
 {
     psa_algorithm_t alg = alg_arg;
     psa_key_derivation_step_t steps[] = {step_arg1, step_arg2, step_arg3};
@@ -4291,6 +4292,10 @@
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     size_t i;
+    psa_key_type_t output_key_type = output_key_type_arg;
+    psa_key_handle_t output_handle = 0;
+    psa_status_t expected_output_status = expected_output_status_arg;
+    psa_status_t actual_output_status;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -4320,10 +4325,29 @@
         }
     }
 
+    if( output_key_type != PSA_KEY_TYPE_NONE )
+    {
+        psa_reset_key_attributes( &attributes );
+        psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
+        psa_set_key_bits( &attributes, 8 );
+        actual_output_status =
+            psa_key_derivation_output_key( &attributes, &operation,
+                                           &output_handle );
+    }
+    else
+    {
+        uint8_t buffer[1];
+        actual_output_status =
+            psa_key_derivation_output_bytes( &operation,
+                                             buffer, sizeof( buffer ) );
+    }
+    TEST_EQUAL( actual_output_status, expected_output_status );
+
 exit:
     psa_key_derivation_abort( &operation );
     for( i = 0; i < ARRAY_LENGTH( handles ); i++ )
         psa_destroy_key( handles[i] );
+    psa_destroy_key( output_handle );
     PSA_DONE( );
 }
 /* END_CASE */