test_suite_pk: replace USE_PSA with CRYPTO_CLIENT in tests with opaque keys

This commit also resolves upcoming issues found in pk_internal.h and
pkwrite.c.

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/library/pk_internal.h b/library/pk_internal.h
index f5924ad..c85e4ff 100644
--- a/library/pk_internal.h
+++ b/library/pk_internal.h
@@ -87,7 +87,7 @@
 {
     mbedtls_ecp_group_id id;
 
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
     if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_OPAQUE) {
         psa_key_attributes_t opaque_attrs = PSA_KEY_ATTRIBUTES_INIT;
         psa_key_type_t opaque_key_type;
@@ -101,7 +101,7 @@
         id = mbedtls_ecc_group_from_psa(curve, psa_get_key_bits(&opaque_attrs));
         psa_reset_key_attributes(&opaque_attrs);
     } else
-#endif /* MBEDTLS_USE_PSA_CRYPTO */
+#endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
     {
 #if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
         id = mbedtls_ecc_group_from_psa(pk->ec_family, pk->ec_bits);
diff --git a/library/pkwrite.c b/library/pkwrite.c
index b9ddcf1..16af2b6 100644
--- a/library/pkwrite.c
+++ b/library/pkwrite.c
@@ -348,7 +348,7 @@
 /******************************************************************************
  * Internal functions for Opaque keys.
  ******************************************************************************/
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
 static int pk_write_opaque_pubkey(unsigned char **p, unsigned char *start,
                                   const mbedtls_pk_context *pk)
 {
@@ -370,7 +370,7 @@
 
     return (int) len;
 }
-#endif /* MBEDTLS_USE_PSA_CRYPTO */
+#endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
 
 /******************************************************************************
  * Generic helpers
@@ -382,7 +382,7 @@
 {
     mbedtls_pk_type_t pk_type = mbedtls_pk_get_type(pk);
 
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
     if (pk_type == MBEDTLS_PK_OPAQUE) {
         psa_key_attributes_t opaque_attrs = PSA_KEY_ATTRIBUTES_INIT;
         psa_key_type_t opaque_key_type;
@@ -424,11 +424,11 @@
         MBEDTLS_ASN1_CHK_ADD(len, pk_write_ec_pubkey(p, start, key));
     } else
 #endif
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
     if (mbedtls_pk_get_type(key) == MBEDTLS_PK_OPAQUE) {
         MBEDTLS_ASN1_CHK_ADD(len, pk_write_opaque_pubkey(p, start, key));
     } else
-#endif /* MBEDTLS_USE_PSA_CRYPTO */
+#endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
     return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
 
     return (int) len;
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index a625b56..a6b4b1b 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -737,27 +737,35 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_FS_IO */
+/* BEGIN_CASE depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_FS_IO:MBEDTLS_PSA_CRYPTO_C */
 void mbedtls_pk_check_pair(char *pub_file, char *prv_file, int ret)
 {
     mbedtls_pk_context pub, prv, alt;
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
     mbedtls_svc_key_id_t opaque_key_id = MBEDTLS_SVC_KEY_ID_INIT;
-#endif /* MBEDTLS_USE_PSA_CRYPTO */
+    int opaque_ret = ret;
+#endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
 
     mbedtls_pk_init(&pub);
     mbedtls_pk_init(&prv);
     mbedtls_pk_init(&alt);
     USE_PSA_INIT();
 
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
-    /* mbedtls_pk_check_pair() returns either PK or ECP error codes depending
-       on MBEDTLS_USE_PSA_CRYPTO so here we dynamically translate between the
-       two */
+    /* In case of EC keys, mbedtls_pk_check_pair() returns either PK or ECP
+     * error codes depending on whether PSA or ECP functions are used to perform
+     * the check.
+     * - For non-opaque keys PSA functions are used when USE_PSA is enabled,
+     *   otherwise legacy ones (ECP) are used.
+     * - For opaque keys PSA functions are always used as soon as opaque keys
+     *   are supported (i.e. MBEDTLS_PSA_CRYPTO_C enabled). */
     if (ret == MBEDTLS_ERR_ECP_BAD_INPUT_DATA) {
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
         ret = MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-    }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
+        opaque_ret = MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+#endif
+    }
 
     TEST_ASSERT(mbedtls_pk_parse_public_keyfile(&pub, pub_file) == 0);
     TEST_ASSERT(mbedtls_pk_parse_keyfile(&prv, prv_file, NULL,
@@ -778,20 +786,23 @@
                     == ret);
     }
 #endif
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
     if (mbedtls_pk_get_type(&prv) == MBEDTLS_PK_ECKEY) {
+        if (ret == MBEDTLS_ERR_ECP_BAD_INPUT_DATA) {
+            ret = MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+        }
         TEST_EQUAL(mbedtls_pk_wrap_as_opaque(&prv, &opaque_key_id,
                                              PSA_ALG_ANY_HASH,
                                              PSA_KEY_USAGE_EXPORT, 0), 0);
         TEST_EQUAL(mbedtls_pk_check_pair(&pub, &prv, mbedtls_test_rnd_std_rand,
-                                         NULL), ret);
+                                         NULL), opaque_ret);
     }
 #endif
 
 exit:
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
     psa_destroy_key(opaque_key_id);
-#endif /* MBEDTLS_USE_PSA_CRYPTO */
+#endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
     mbedtls_pk_free(&pub);
     mbedtls_pk_free(&prv);
     mbedtls_pk_free(&alt);
@@ -1250,7 +1261,7 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_RSA_C:MBEDTLS_USE_PSA_CRYPTO */
+/* BEGIN_CASE depends_on:MBEDTLS_RSA_C:MBEDTLS_PSA_CRYPTO_C */
 void pk_wrap_rsa_decrypt_test_vec(data_t *cipher, int mod,
                                   char *input_P, char *input_Q,
                                   char *input_N, char *input_E,
@@ -1489,7 +1500,7 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_MD_CAN_SHA256:MBEDTLS_USE_PSA_CRYPTO:MBEDTLS_TEST_PK_PSA_SIGN */
+/* BEGIN_CASE depends_on:MBEDTLS_MD_CAN_SHA256:MBEDTLS_PSA_CRYPTO_C:MBEDTLS_TEST_PK_PSA_SIGN */
 void pk_psa_sign(int curve_or_keybits, int psa_type, int expected_bits)
 {
     mbedtls_pk_context pk;