mbedtls_pk_get_psa_attributes: support MBEDTLS_PK_USE_PSA_EC_DATA

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 930d015..e522ea7 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -111,7 +111,14 @@
         mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECDSA) {
         int ret;
 
-#if defined(MBEDTLS_ECP_C)
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+        ret = pk_genkey_ec(pk, curve_or_keybits);
+        if (ret != 0) {
+            return ret;
+        }
+
+        return 0;
+#else
         ret = mbedtls_ecp_group_load(&mbedtls_pk_ec_rw(*pk)->grp, curve_or_keybits);
         if (ret != 0) {
             return ret;
@@ -120,16 +127,7 @@
                                        &mbedtls_pk_ec_rw(*pk)->d,
                                        &mbedtls_pk_ec_rw(*pk)->Q,
                                        mbedtls_test_rnd_std_rand, NULL);
-#endif /* MBEDTLS_ECP_C */
-
-#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-        ret = pk_genkey_ec(pk, curve_or_keybits);
-        if (ret != 0) {
-            return ret;
-        }
-
-        return 0;
-#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+#endif /* MBEDTLS_ECP_C && !MBEDTLS_PK_USE_PSA_EC_DATA */
 
     }
 #endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
@@ -208,8 +206,24 @@
             *psa_type = PSA_KEY_TYPE_ECC_KEY_PAIR(mbedtls_ecc_group_to_psa(grp_id, &bits));
             TEST_EQUAL(pk_genkey(pk, grp_id), 0);
             if (!want_pair) {
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+                psa_key_attributes_t pub_attributes = PSA_KEY_ATTRIBUTES_INIT;
+                psa_set_key_type(&pub_attributes,
+                                 PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(*psa_type));
+                psa_set_key_usage_flags(&pub_attributes,
+                                        PSA_KEY_USAGE_EXPORT |
+                                        PSA_KEY_USAGE_COPY |
+                                        PSA_KEY_USAGE_VERIFY_MESSAGE |
+                                        PSA_KEY_USAGE_VERIFY_HASH);
+                psa_set_key_algorithm(&pub_attributes, PSA_ALG_ECDSA_ANY);
+                PSA_ASSERT(psa_destroy_key(pk->priv_id));
+                PSA_ASSERT(psa_import_key(&pub_attributes,
+                                          pk->pub_raw, pk->pub_raw_len,
+                                          &pk->priv_id));
+#else
                 mbedtls_ecp_keypair *ec = mbedtls_pk_ec(*pk);
                 mbedtls_mpi_free(&ec->d);
+#endif
             }
             break;
         }