diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index fc5684f..653467b 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -68,6 +68,7 @@
     size_t reexported_length;
     psa_key_type_t got_type;
     size_t got_bits;
+    psa_key_policy_t policy = {0};
 
     data = unhexify_alloc( hex, &data_size );
     TEST_ASSERT( data != NULL );
@@ -81,6 +82,13 @@
     }
     TEST_ASSERT( psa_crypto_init( ) == PSA_SUCCESS );
 
+    psa_key_policy_init( &policy );
+
+    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_EXPORT, 
+                              PSA_ALG_VENDOR_FLAG );
+
+    TEST_ASSERT( psa_set_key_policy( slot, &policy ) == PSA_SUCCESS );
+
     /* Import the key */
     TEST_ASSERT( psa_import_key( slot, type,
                                  data, data_size ) == PSA_SUCCESS );
@@ -107,6 +115,8 @@
     }
     else
     {
+        TEST_ASSERT( psa_set_key_policy( slot2, &policy ) == PSA_SUCCESS );
+
         TEST_ASSERT( psa_import_key( slot2, type,
                                      exported, export_size ) ==
                      PSA_SUCCESS );
@@ -276,6 +286,7 @@
     unsigned char *signature = NULL;
     size_t signature_size;
     size_t signature_length = 0xdeadbeef;
+    psa_key_policy_t policy = {0};
 
     key_data = unhexify_alloc( key_hex, &key_size );
     TEST_ASSERT( key_data != NULL );
@@ -286,6 +297,12 @@
 
     TEST_ASSERT( psa_crypto_init( ) == PSA_SUCCESS );
 
+    psa_key_policy_init( &policy );
+
+    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_SIGN, alg_arg );
+
+    TEST_ASSERT( psa_set_key_policy( slot, &policy ) == PSA_SUCCESS );
+
     TEST_ASSERT( psa_import_key( slot, key_type,
                                  key_data, key_size ) == PSA_SUCCESS );
     TEST_ASSERT( psa_get_key_information( slot,
@@ -331,6 +348,7 @@
     psa_status_t expected_status = expected_status_arg;
     unsigned char *signature = NULL;
     size_t signature_length = 0xdeadbeef;
+    psa_key_policy_t policy = {0};
 
     key_data = unhexify_alloc( key_hex, &key_size );
     TEST_ASSERT( key_data != NULL );
@@ -341,6 +359,12 @@
 
     TEST_ASSERT( psa_crypto_init( ) == PSA_SUCCESS );
 
+    psa_key_policy_init( &policy );
+
+    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_SIGN, alg_arg );
+
+    TEST_ASSERT( psa_set_key_policy( slot, &policy ) == PSA_SUCCESS );
+
     TEST_ASSERT( psa_import_key( slot, key_type,
                                  key_data, key_size ) == PSA_SUCCESS );