diff --git a/tests/include/test/psa_test_wrappers.h b/tests/include/test/psa_test_wrappers.h
index 880f34c7af..904baad7d8 100644
--- a/tests/include/test/psa_test_wrappers.h
+++ b/tests/include/test/psa_test_wrappers.h
@@ -363,7 +363,7 @@ psa_status_t mbedtls_test_wrap_psa_export_public_key_iop_complete(
 
 psa_status_t mbedtls_test_wrap_psa_export_public_key_iop_setup(
     psa_export_public_key_iop_t *arg0_operation,
-    psa_key_id_t arg1_key);
+    mbedtls_svc_key_id_t arg1_key);
 #define psa_export_public_key_iop_setup(arg0_operation, arg1_key) \
     mbedtls_test_wrap_psa_export_public_key_iop_setup(arg0_operation, arg1_key)
 
diff --git a/tests/src/psa_test_wrappers.c b/tests/src/psa_test_wrappers.c
index 65d81573a7..3df07723a2 100644
--- a/tests/src/psa_test_wrappers.c
+++ b/tests/src/psa_test_wrappers.c
@@ -623,7 +623,7 @@ psa_status_t mbedtls_test_wrap_psa_export_public_key_iop_complete(
 /* Wrapper for psa_export_public_key_iop_setup */
 psa_status_t mbedtls_test_wrap_psa_export_public_key_iop_setup(
     psa_export_public_key_iop_t *arg0_operation,
-    psa_key_id_t arg1_key)
+    mbedtls_svc_key_id_t arg1_key)
 {
     psa_status_t status = (psa_export_public_key_iop_setup)(arg0_operation, arg1_key);
     return status;
diff --git a/tf-psa-crypto/include/psa/crypto.h b/tf-psa-crypto/include/psa/crypto.h
index c1ef041f77..6ea945d3ec 100644
--- a/tf-psa-crypto/include/psa/crypto.h
+++ b/tf-psa-crypto/include/psa/crypto.h
@@ -5677,7 +5677,7 @@ uint32_t psa_export_public_key_iop_get_num_ops(psa_export_public_key_iop_t *oper
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY \emptydescription
  */
 psa_status_t psa_export_public_key_iop_setup(psa_export_public_key_iop_t *operation,
-                                             psa_key_id_t key);
+                                             mbedtls_svc_key_id_t key);
 
 /**
  * \brief                       Continue and eventually complete the action of