diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index e1a8e1c..8f5e7f6 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1822,13 +1822,15 @@
     mbedtls_pk_context pk;
     unsigned char hash[32];
     unsigned char sig[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
-    unsigned char pkey_legacy[200];
-    unsigned char pkey_psa[200];
-    unsigned char *pkey_legacy_start, *pkey_psa_start;
-    size_t sig_len, klen_legacy, klen_psa;
+    unsigned char legacy_pub_key[PSA_EXPORT_PUBLIC_KEY_MAX_SIZE];
+    unsigned char opaque_pub_key[PSA_EXPORT_PUBLIC_KEY_MAX_SIZE];
+    size_t sig_len, legacy_pub_key_len, opaque_pub_key_len;
     int ret;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+#if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
+    mbedtls_ecp_group_id ecp_grp_id;
+#endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
 
     /*
      * Following checks are perfomed:
@@ -1853,9 +1855,9 @@
 #endif /* MBEDTLS_RSA_C && MBEDTLS_GENPRIME */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
-        mbedtls_ecp_group_id grpid = mbedtls_ecc_group_from_psa(psa_type, bits);
+        ecp_grp_id = mbedtls_ecc_group_from_psa(psa_type, bits);
         TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)) == 0);
-        TEST_ASSERT(pk_genkey(&pk, grpid) == 0);
+        TEST_ASSERT(pk_genkey(&pk, ecp_grp_id) == 0);
     } else
 #endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
     {
@@ -1864,19 +1866,21 @@
 
     /* Export public key from the non-opaque PK context we just created. */
 #if defined(MBEDTLS_PK_WRITE_C)
-    ret = mbedtls_pk_write_pubkey_der(&pk, pkey_legacy, sizeof(pkey_legacy));
+    ret = mbedtls_pk_write_pubkey_der(&pk, legacy_pub_key, sizeof(legacy_pub_key));
     TEST_ASSERT(ret >= 0);
-    klen_legacy = (size_t) ret;
-    /* mbedtls_pk_write_pubkey_der() writes backwards in the data buffer. */
-    pkey_legacy_start = pkey_legacy + sizeof(pkey_legacy) - klen_legacy;
+    legacy_pub_key_len = (size_t) ret;
+    /* mbedtls_pk_write_pubkey_der() writes backwards in the data buffer so we
+     * shift data back to the beginning of the buffer. */
+    memmove(legacy_pub_key,
+            legacy_pub_key + sizeof(legacy_pub_key) - legacy_pub_key_len,
+            legacy_pub_key_len);
 #else /* MBEDTLS_PK_WRITE_C */
     ret = mbedtls_ecp_point_write_binary(&(mbedtls_pk_ec_ro(pk)->grp),
                                          &(mbedtls_pk_ec_ro(pk)->Q),
                                          MBEDTLS_ECP_PF_UNCOMPRESSED,
-                                         &klen_legacy, pkey_legacy,
-                                         sizeof(pkey_legacy));
+                                         &legacy_pub_key_len, legacy_pub_key,
+                                         sizeof(legacy_pub_key));
     TEST_EQUAL(ret, 0);
-    pkey_legacy_start = pkey_legacy;
 #endif /* MBEDTLS_PK_WRITE_C */
 
     /* Turn the PK context into an opaque one. */
@@ -1905,26 +1909,23 @@
 
     /* Export public key from the opaque PK context. */
 #if defined(MBEDTLS_PK_WRITE_C)
-    ret = mbedtls_pk_write_pubkey_der(&pk, pkey_psa,
-                                      sizeof(pkey_psa));
+    ret = mbedtls_pk_write_pubkey_der(&pk, opaque_pub_key, sizeof(opaque_pub_key));
     TEST_ASSERT(ret >= 0);
-    klen_psa = (size_t) ret;
+    opaque_pub_key_len = (size_t) ret;
     /* mbedtls_pk_write_pubkey_der() writes backwards in the data buffer. */
-    pkey_psa_start = pkey_psa + sizeof(pkey_psa) - klen_psa;
+    memmove(opaque_pub_key,
+            opaque_pub_key + sizeof(opaque_pub_key) - opaque_pub_key_len,
+            opaque_pub_key_len);
 #else
-    psa_status_t status;
-
-    status = psa_export_public_key(key_id, pkey_psa, sizeof(pkey_psa),
-                                   &klen_psa);
-    TEST_EQUAL(status, PSA_SUCCESS);
-    pkey_psa_start = pkey_psa;
+    TEST_EQUAL(psa_export_public_key(key_id, opaque_pub_key, sizeof(opaque_pub_key),
+                                     &opaque_pub_key_len), PSA_SUCCESS);
 #endif /* MBEDTLS_PK_WRITE_C */
 
     /* Check that the public keys of opaque and non-opaque PK contexts match. */
-    TEST_ASSERT(klen_psa == klen_legacy);
-    TEST_ASSERT(memcmp(pkey_psa_start, pkey_legacy_start, klen_psa) == 0);
+    TEST_EQUAL(opaque_pub_key_len, legacy_pub_key_len);
+    TEST_MEMORY_COMPARE(opaque_pub_key, opaque_pub_key_len, legacy_pub_key, legacy_pub_key_len);
 
-    /* Destroy the opaque PK context. */
+    /* Destroy the opaque PK context and the wrapped PSA key. */
     mbedtls_pk_free(&pk);
     TEST_ASSERT(PSA_SUCCESS == psa_destroy_key(key_id));
 
@@ -1936,17 +1937,13 @@
      * on "PK_WRITE" and "PK_PARSE" if required */
     mbedtls_pk_init(&pk);
 #if defined(MBEDTLS_PK_WRITE_C) && defined(MBEDTLS_PK_PARSE_C)
-    TEST_EQUAL(mbedtls_pk_parse_public_key(&pk, pkey_legacy_start,
-                                           klen_legacy), 0);
+    TEST_EQUAL(mbedtls_pk_parse_public_key(&pk, legacy_pub_key, legacy_pub_key_len), 0);
 #else
-    TEST_EQUAL(mbedtls_pk_setup(&pk,
-                                mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)), 0);
-    TEST_EQUAL(mbedtls_ecp_group_load(
-                   &(mbedtls_pk_ec_rw(pk)->grp),
-                   (mbedtls_ecp_group_id) curve_or_keybits), 0);
+    TEST_EQUAL(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)), 0);
+    TEST_EQUAL(mbedtls_ecp_group_load(&(mbedtls_pk_ec_rw(pk)->grp), ecp_grp_id), 0);
     TEST_EQUAL(mbedtls_ecp_point_read_binary(&(mbedtls_pk_ec_ro(pk)->grp),
                                              &(mbedtls_pk_ec_rw(pk)->Q),
-                                             pkey_legacy_start, klen_legacy), 0);
+                                             legacy_pub_key, legacy_pub_key_len), 0);
 #endif
     TEST_ASSERT(mbedtls_pk_verify(&pk, MBEDTLS_MD_SHA256,
                                   hash, sizeof(hash), sig, sig_len) == 0);