diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index e0bedf7..436ed7c 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -342,8 +342,17 @@
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_RSA_C
 import_rsa_made_up:PSA_VENDOR_RSA_MAX_KEY_BITS+8:0:PSA_ERROR_NOT_SUPPORTED
 
-PSA key policy set and get
-key_policy:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CBC_NO_PADDING
+PSA key policy: AES
+depends_on:MBEDTLS_AES_C
+check_key_policy:PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_ENCRYPT:PSA_ALG_CBC_NO_PADDING
+
+PSA key policy: ECC SECP256R1, sign
+depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+check_key_policy:PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_CURVE_SECP256R1 ):256:PSA_KEY_USAGE_SIGN_HASH:PSA_ALG_ECDSA_ANY
+
+PSA key policy: ECC SECP256R1, sign+verify
+depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+check_key_policy:PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_CURVE_SECP256R1 ):256:PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY
 
 Key attributes initializers zero properly
 key_attributes_init:
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 3ce8df8..d62d3c1 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1635,27 +1635,29 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void key_policy( int usage_arg, int alg_arg )
+void check_key_policy( int type_arg, int bits_arg,
+                       int usage_arg, int alg_arg )
 {
     psa_key_handle_t handle = 0;
+    psa_key_type_t key_type = type_arg;
+    size_t bits = bits_arg;
     psa_algorithm_t alg = alg_arg;
     psa_key_usage_t usage = usage_arg;
-    psa_key_type_t key_type = PSA_KEY_TYPE_AES;
-    unsigned char key[32] = {0};
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
-    memset( key, 0x2a, sizeof( key ) );
-
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_usage_flags( &attributes, usage );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, key_type );
+    psa_set_key_bits( &attributes, bits );
 
-    PSA_ASSERT( psa_import_key( &attributes, key, sizeof( key ), &handle ) );
+    PSA_ASSERT( psa_generate_key( &attributes, &handle ) );
+    psa_reset_key_attributes( &attributes );
 
     PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
     TEST_EQUAL( psa_get_key_type( &attributes ), key_type );
+    TEST_EQUAL( psa_get_key_bits( &attributes ), bits );
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), usage );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), alg );