diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 6cd6bee9e..7b9bbfe54 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -7087,9 +7087,12 @@ void derive_output( int alg_arg,
                     goto exit;
                 break;
             default:
-                PSA_ASSERT( psa_key_derivation_input_bytes(
+                TEST_EQUAL( psa_key_derivation_input_bytes(
                                 &operation, steps[i],
-                                inputs[i]->x, inputs[i]->len ) );
+                                inputs[i]->x, inputs[i]->len ), statuses[i] );
+
+                if( statuses[i] != PSA_SUCCESS )
+                    goto exit;
                 break;
         }
     }
@@ -7102,6 +7105,10 @@ void derive_output( int alg_arg,
     if( derive_type == 1 ) // output key
     {
         psa_status_t expected_status = PSA_ERROR_NOT_PERMITTED;
+        size_t bits = 48; // default for Mix-PSK-to-MS
+
+        if( PSA_ALG_IS_HKDF_EXTRACT( alg ) || PSA_ALG_IS_HKDF_EXPAND( alg ))
+            bits = PSA_HASH_LENGTH( alg );
 
         /* For output key derivation secret must be provided using
            input key, otherwise operation is not permitted. */
@@ -7111,7 +7118,7 @@ void derive_output( int alg_arg,
         psa_set_key_usage_flags( &attributes4, PSA_KEY_USAGE_EXPORT );
         psa_set_key_algorithm( &attributes4, alg );
         psa_set_key_type( &attributes4, PSA_KEY_TYPE_DERIVE );
-        psa_set_key_bits( &attributes4, 48 );
+        psa_set_key_bits( &attributes4, bits );
 
         TEST_EQUAL( psa_key_derivation_output_key( &attributes4, &operation,
                                         &derived_key ), expected_status );