Update test for Opaque PK key

Signed-off-by: Neil Armstrong <narmstrong@baylibre.com>
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index ccb90e7..0a64bc9 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -88,11 +88,11 @@
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 
 /*
- * Generate a key using PSA and return the key identifier of that key,
+ * Generate an ECC key using PSA and return the key identifier of that key,
  * or 0 if the key generation failed.
  * The key uses NIST P-256 and is usable for signing with SHA-256.
  */
-mbedtls_svc_key_id_t pk_psa_genkey( void )
+mbedtls_svc_key_id_t pk_psa_genkey_ecc( void )
 {
     mbedtls_svc_key_id_t key;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -109,6 +109,27 @@
 exit:
     return( key );
 }
+
+/*
+ * Generate an RSA key using PSA and return the key identifier of that key,
+ * or 0 if the key generation failed.
+ */
+mbedtls_svc_key_id_t pk_psa_genkey_rsa( void )
+{
+    mbedtls_svc_key_id_t key;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    const psa_key_type_t type = PSA_KEY_TYPE_RSA_KEY_PAIR;
+    const size_t bits = 256;
+
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_HASH );
+    psa_set_key_algorithm( &attributes, PSA_ALG_RSA_PKCS1V15_SIGN_RAW );
+    psa_set_key_type( &attributes, type );
+    psa_set_key_bits( &attributes, bits );
+    PSA_ASSERT( psa_generate_key( &attributes, &key ) );
+
+exit:
+    return( key );
+}
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 /* END_HEADER */
 
@@ -117,8 +138,8 @@
  * END_DEPENDENCIES
  */
 
-/* BEGIN_CASE depends_on:MBEDTLS_USE_PSA_CRYPTO:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED */
-void pk_psa_utils(  )
+/* BEGIN_CASE depends_on:MBEDTLS_USE_PSA_CRYPTO */
+void pk_psa_utils( int key_is_rsa )
 {
     mbedtls_pk_context pk, pk2;
     mbedtls_svc_key_id_t key;
@@ -145,7 +166,10 @@
     mbedtls_pk_free( &pk );
     mbedtls_pk_init( &pk );
 
-    key = pk_psa_genkey();
+    if( key_is_rsa )
+        key = pk_psa_genkey_rsa();
+    else
+        key = pk_psa_genkey_ecc();
     if( mbedtls_svc_key_id_is_null( key ) )
         goto exit;
 
@@ -157,9 +181,18 @@
     TEST_ASSERT( mbedtls_pk_get_bitlen( &pk ) == bitlen );
     TEST_ASSERT( mbedtls_pk_get_len( &pk ) == bitlen / 8 );
 
-    TEST_ASSERT( mbedtls_pk_can_do( &pk, MBEDTLS_PK_ECKEY ) == 1 );
-    TEST_ASSERT( mbedtls_pk_can_do( &pk, MBEDTLS_PK_ECDSA ) == 1 );
-    TEST_ASSERT( mbedtls_pk_can_do( &pk, MBEDTLS_PK_RSA ) == 0 );
+    if( key_is_rsa )
+    {
+        TEST_ASSERT( mbedtls_pk_can_do( &pk, MBEDTLS_PK_ECKEY ) == 0 );
+        TEST_ASSERT( mbedtls_pk_can_do( &pk, MBEDTLS_PK_ECDSA ) == 0 );
+        TEST_ASSERT( mbedtls_pk_can_do( &pk, MBEDTLS_PK_RSA ) == 1 );
+    }
+    else
+    {
+        TEST_ASSERT( mbedtls_pk_can_do( &pk, MBEDTLS_PK_ECKEY ) == 1 );
+        TEST_ASSERT( mbedtls_pk_can_do( &pk, MBEDTLS_PK_ECDSA ) == 1 );
+        TEST_ASSERT( mbedtls_pk_can_do( &pk, MBEDTLS_PK_RSA ) == 0 );
+    }
 
     /* unsupported operations: verify, decrypt, encrypt */
     TEST_ASSERT( mbedtls_pk_verify( &pk, md_alg,
@@ -175,8 +208,12 @@
                  == MBEDTLS_ERR_PK_TYPE_MISMATCH );
 
     /* unsupported functions: check_pair, debug */
-    TEST_ASSERT( mbedtls_pk_setup( &pk2,
-                 mbedtls_pk_info_from_type( MBEDTLS_PK_ECKEY ) ) == 0 );
+    if( key_is_rsa )
+        TEST_ASSERT( mbedtls_pk_setup( &pk2,
+                     mbedtls_pk_info_from_type( MBEDTLS_PK_RSA ) ) == 0 );
+    else
+        TEST_ASSERT( mbedtls_pk_setup( &pk2,
+                     mbedtls_pk_info_from_type( MBEDTLS_PK_ECKEY ) ) == 0 );
     TEST_ASSERT( mbedtls_pk_check_pair( &pk, &pk2,
                                         mbedtls_test_rnd_std_rand, NULL )
                  == MBEDTLS_ERR_PK_TYPE_MISMATCH );