Update tests for extended key usage policies

Signed-off-by: gabor-mezei-arm <gabor.mezei@arm.com>
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 41e2972..3e193c5 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -322,7 +322,8 @@
 
     PSA_ASSERT( psa_get_key_attributes( key, &got_attributes ) );
     TEST_EQUAL( psa_get_key_type( &got_attributes ), type );
-    TEST_EQUAL( psa_get_key_usage_flags( &got_attributes ), usage );
+    TEST_EQUAL( psa_get_key_usage_flags( &got_attributes ),
+                                         update_key_usage_flags( usage ) );
     TEST_EQUAL( psa_get_key_algorithm( &got_attributes ), alg );
     ASSERT_NO_SLOT_NUMBER( &got_attributes );
 
@@ -1162,6 +1163,8 @@
     PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                 &key ) );
 
+    /* Update the usage flags to obtain extended usage flags */
+    usage = update_key_usage_flags( usage );
     PSA_ASSERT( psa_get_key_attributes( key, &got_attributes ) );
     TEST_EQUAL( psa_get_key_usage_flags( &got_attributes ), usage );
     TEST_EQUAL( psa_get_key_algorithm( &got_attributes ), alg );
@@ -5186,7 +5189,8 @@
                 PSA_KEY_LIFETIME_PERSISTENT );
     TEST_EQUAL( psa_get_key_type( &attributes ), type );
     TEST_EQUAL( psa_get_key_bits( &attributes ), bits );
-    TEST_EQUAL( psa_get_key_usage_flags( &attributes ), usage_flags );
+    TEST_EQUAL( psa_get_key_usage_flags( &attributes ),
+                                         update_key_usage_flags( usage_flags ) );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), alg );
 
     /* Export the key again if permitted by the key policy. */