HKDF: positive tests
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 278554f..cf72e48 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -8,6 +8,8 @@
 #include "mbedtls/asn1write.h"
 #include "psa/crypto.h"
 
+#define ARRAY_LENGTH( array ) ( sizeof( array ) / sizeof( *( array ) ) )
+
 #if(UINT32_MAX > SIZE_MAX)
 #define PSA_CRYPTO_TEST_SIZE_T_RANGE( x ) ( ( x ) <= SIZE_MAX )
 #else
@@ -2484,6 +2486,104 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
+void derive_output( int alg_arg,
+                    data_t *key_data,
+                    data_t *salt,
+                    data_t *label,
+                    int requested_capacity_arg,
+                    data_t *expected_output1,
+                    data_t *expected_output2 )
+{
+    psa_key_slot_t slot = 1;
+    psa_algorithm_t alg = alg_arg;
+    size_t requested_capacity = requested_capacity_arg;
+    psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
+    uint8_t *expected_outputs[2] =
+        {expected_output1->x, expected_output2->x};
+    size_t output_sizes[2] =
+        {expected_output1->len, expected_output2->len};
+    size_t output_buffer_size = 0;
+    uint8_t *output_buffer = NULL;
+    size_t expected_capacity;
+    size_t current_capacity;
+    psa_key_policy_t policy;
+    psa_status_t status;
+    unsigned i;
+
+    for( i = 0; i < ARRAY_LENGTH( expected_outputs ); i++ )
+    {
+        if( output_sizes[i] > output_buffer_size )
+            output_buffer_size = output_sizes[i];
+        if( output_sizes[i] == 0 )
+            expected_outputs[i] = NULL;
+    }
+    output_buffer = mbedtls_calloc( 1, output_buffer_size );
+    TEST_ASSERT( output_buffer != NULL );
+    TEST_ASSERT( psa_crypto_init( ) == PSA_SUCCESS );
+
+    psa_key_policy_init( &policy );
+    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
+    TEST_ASSERT( psa_set_key_policy( slot, &policy ) == PSA_SUCCESS );
+
+    TEST_ASSERT( psa_import_key( slot, PSA_KEY_TYPE_DERIVE,
+                                 key_data->x,
+                                 key_data->len ) == PSA_SUCCESS );
+
+    /* Extraction phase. */
+    TEST_ASSERT( psa_key_derivation( &generator, slot, alg,
+                                     salt->x, salt->len,
+                                     label->x, label->len,
+                                     requested_capacity ) == PSA_SUCCESS );
+    TEST_ASSERT( psa_get_generator_capacity( &generator,
+                                             &current_capacity ) ==
+                 PSA_SUCCESS );
+    TEST_ASSERT( current_capacity == requested_capacity );
+    expected_capacity = requested_capacity;
+
+    /* Expansion phase. */
+    for( i = 0; i < ARRAY_LENGTH( expected_outputs ); i++ )
+    {
+        /* Read some bytes. */
+        status = psa_generator_read( &generator,
+                                     output_buffer, output_sizes[i] );
+        if( expected_capacity == 0 && output_sizes[i] == 0 )
+        {
+            /* Reading 0 bytes when 0 bytes are available can go either way. */
+            TEST_ASSERT( status == PSA_SUCCESS ||
+                         status == PSA_ERROR_INSUFFICIENT_CAPACITY );
+            continue;
+        }
+        else if( expected_capacity == 0 ||
+                 output_sizes[i] > expected_capacity )
+        {
+            /* Capacity exceeded. */
+            TEST_ASSERT( status == PSA_ERROR_INSUFFICIENT_CAPACITY );
+            expected_capacity = 0;
+            continue;
+        }
+        /* Success. Check the read data. */
+        TEST_ASSERT( status == PSA_SUCCESS );
+        if( output_sizes[i] != 0 )
+            TEST_ASSERT( memcmp( output_buffer, expected_outputs[i],
+                                 output_sizes[i] ) == 0 );
+        /* Check the generator status. */
+        expected_capacity -= output_sizes[i];
+        TEST_ASSERT( psa_get_generator_capacity( &generator,
+                                                 &current_capacity ) ==
+                     PSA_SUCCESS );
+        TEST_ASSERT( expected_capacity == current_capacity );
+    }
+    TEST_ASSERT( psa_generator_abort( &generator ) == PSA_SUCCESS );
+
+exit:
+    mbedtls_free( output_buffer );
+    psa_generator_abort( &generator );
+    psa_destroy_key( slot );
+    mbedtls_psa_crypto_free( );
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
 void generate_random( int bytes_arg )
 {
     size_t bytes = bytes_arg;