diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 3b0e18936..d1fc2582d 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -294,16 +294,17 @@ exit:
     ((void) 0)
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
-uint64_t parse_hex_string(data_t *hex_string)
+#define INPUT_INTEGER 0x10000   /* Out of range of psa_key_type_t */
+
+uint64_t parse_binary_string(data_t *bin_string)
 {
     uint64_t result = 0;
-    if (hex_string->len > 8) {
-        return 0;
+    TEST_LE_U(bin_string->len, 8);
+    for (size_t i = 0; i < bin_string->len; i++) {
+        result |= ((bin_string->x)[i]) << (i * 8);
     }
-    for (size_t i = 0; i < hex_string->len; i++) {
-        result |= ((hex_string->x)[i]) << (i * 8);
-    }
-    return result;
+exit:
+    return result; /* returns 0 if len > 8 */
 }
 
 /* An overapproximation of the amount of storage needed for a key of the
@@ -330,10 +331,6 @@ typedef enum {
     USE_GIVEN_TAG = 1,
 } tag_usage_method_t;
 
-typedef enum {
-    INPUT_BYTES = 0,
-    INPUT_INTEGER = 1
-} key_derivation_input_method_t;
 
 /*!
  * \brief                           Internal Function for AEAD multipart tests.
@@ -8464,19 +8461,17 @@ exit:
 
 /* BEGIN_CASE */
 void derive_input(int alg_arg,
-                  int step_arg1, int key_type_arg1, int input_type_arg1,
-                  data_t *input1, int expected_status_arg1,
-                  int step_arg2, int key_type_arg2, int input_type_arg2,
-                  data_t *input2, int expected_status_arg2,
-                  int step_arg3, int key_type_arg3, int input_type_arg3,
-                  data_t *input3, int expected_status_arg3,
+                  int step_arg1, int key_type_arg1, data_t *input1,
+                  int expected_status_arg1,
+                  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 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 };
-    psa_key_type_t key_types[] = { key_type_arg1, key_type_arg2, key_type_arg3 };
-    key_derivation_input_method_t input_types[] =
-    { input_type_arg1, input_type_arg2, input_type_arg3 };
+    uint32_t key_types[] = { key_type_arg1, key_type_arg2, key_type_arg3 };
     psa_status_t expected_statuses[] = { expected_status_arg1,
                                          expected_status_arg2,
                                          expected_status_arg3 };
@@ -8503,12 +8498,13 @@ void derive_input(int alg_arg,
         mbedtls_test_set_step(i);
         if (steps[i] == 0) {
             /* Skip this step */
-        } else if (key_types[i] != PSA_KEY_TYPE_NONE) {
-            psa_set_key_type(&attributes, key_types[i]);
+        } else if (((psa_key_type_t) key_types[i]) != PSA_KEY_TYPE_NONE &&
+                   key_types[i] != INPUT_INTEGER) {
+            psa_set_key_type(&attributes, ((psa_key_type_t) key_types[i]));
             PSA_ASSERT(psa_import_key(&attributes,
                                       inputs[i]->x, inputs[i]->len,
                                       &keys[i]));
-            if (PSA_KEY_TYPE_IS_KEY_PAIR(key_types[i]) &&
+            if (PSA_KEY_TYPE_IS_KEY_PAIR((psa_key_type_t) key_types[i]) &&
                 steps[i] == PSA_KEY_DERIVATION_INPUT_SECRET) {
                 // When taking a private key as secret input, use key agreement
                 // to add the shared secret to the derivation
@@ -8521,16 +8517,16 @@ void derive_input(int alg_arg,
                            expected_statuses[i]);
             }
         } else {
-            if (input_types[i] == INPUT_BYTES) {
+            if (key_types[i] == INPUT_INTEGER) {
+                TEST_EQUAL(psa_key_derivation_input_integer(
+                               &operation, steps[i],
+                               parse_binary_string(inputs[i])),
+                           expected_statuses[i]);
+            } else {
                 TEST_EQUAL(psa_key_derivation_input_bytes(
                                &operation, steps[i],
                                inputs[i]->x, inputs[i]->len),
                            expected_statuses[i]);
-            } else if (input_types[i] == INPUT_INTEGER) {
-                TEST_EQUAL(psa_key_derivation_input_integer(
-                               &operation, steps[i],
-                               parse_hex_string(inputs[i])),
-                           expected_statuses[i]);
             }
         }
     }