diff --git a/library/pkwrite.h b/library/pkwrite.h
index 544ab2f329..01dc3d2f0f 100644
--- a/library/pkwrite.h
+++ b/library/pkwrite.h
@@ -109,4 +109,13 @@
 #define MBEDTLS_PK_ECP_PRV_DER_MAX_BYTES   0
 
 #endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+
+/* Define the maximum available public key DER length based on the supported
+ * key types (EC and/or RSA). */
+#if (MBEDTLS_PK_RSA_PUB_DER_MAX_BYTES > MBEDTLS_PK_ECP_PUB_DER_MAX_BYTES)
+#define MBEDTLS_PK_WRITE_PUBKEY_MAX_SIZE    MBEDTLS_PK_RSA_PUB_DER_MAX_BYTES
+#else
+#define MBEDTLS_PK_WRITE_PUBKEY_MAX_SIZE    MBEDTLS_PK_ECP_PUB_DER_MAX_BYTES
+#endif
+
 #endif /* MBEDTLS_PK_WRITE_H */
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 272e9f67a9..d116050bf6 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -23,6 +23,9 @@
 
 #include <test/psa_exercise_key.h>
 
+/* Needed for the definition of MBEDTLS_PK_WRITE_PUBKEY_MAX_SIZE. */
+#include "pkwrite.h"
+
 /* Used for properly sizing the key buffer in pk_genkey_ec() */
 #include "psa_util_internal.h"
 
@@ -1828,13 +1831,8 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     mbedtls_pk_context pk;
     unsigned char hash[32];
     unsigned char sig[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
-    /* These buffers will be used to contain the key in DER format. Unfortunately
-     * when only EC is supported on the PSA side (i.e. no RSA or DH)
-     * PSA_EXPORT_PUBLIC_KEY_MAX_SIZE falls to PSA_KEY_EXPORT_ECC_PUBLIC_KEY_MAX_SIZE(),
-     * but this is not enough to contain the DER representation of the public key.
-     * Therefore we pick the RSA size to be safe. */
-    unsigned char legacy_pub_key[RSA_WRITE_PUBKEY_MAX_SIZE];
-    unsigned char opaque_pub_key[RSA_WRITE_PUBKEY_MAX_SIZE];
+    unsigned char legacy_pub_key[MBEDTLS_PK_WRITE_PUBKEY_MAX_SIZE];
+    unsigned char opaque_pub_key[MBEDTLS_PK_WRITE_PUBKEY_MAX_SIZE];
     size_t sig_len, legacy_pub_key_len, opaque_pub_key_len;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;