diff --git a/scripts/mbedtls_dev/psa_storage.py b/scripts/mbedtls_dev/psa_storage.py
index cde6c80..0a8a0a3 100644
--- a/scripts/mbedtls_dev/psa_storage.py
+++ b/scripts/mbedtls_dev/psa_storage.py
@@ -102,9 +102,9 @@
     """The latest version of the storage format."""
 
     IMPLICIT_USAGE_FLAGS = {
-        Expr('PSA_KEY_USAGE_SIGN_HASH'): Expr('PSA_KEY_USAGE_SIGN_MESSAGE'),
-        Expr('PSA_KEY_USAGE_VERIFY_HASH'): Expr('PSA_KEY_USAGE_VERIFY_MESSAGE')
-    } #type: Dict[Expr, Expr]
+        'PSA_KEY_USAGE_SIGN_HASH': 'PSA_KEY_USAGE_SIGN_MESSAGE',
+        'PSA_KEY_USAGE_VERIFY_HASH': 'PSA_KEY_USAGE_VERIFY_MESSAGE'
+    } #type: Dict[str, str]
     """Mapping of usage flags to the flags that they imply."""
 
     IMPLICIT_USAGE_FLAGS_KEY_RESTRICTION = {
@@ -138,10 +138,10 @@
 
         if usage_extension:
             for flag, extension in self.IMPLICIT_USAGE_FLAGS.items():
-                if self.original_usage.value() & flag.value() and \
-                   self.original_usage.value() & extension.value() == 0:
+                if self.original_usage.value() & Expr(flag).value() and \
+                   self.original_usage.value() & Expr(extension).value() == 0:
                     self.updated_usage = Expr(self.updated_usage.string +
-                                              ' | ' + extension.string)
+                                              ' | ' + extension)
 
     MAGIC = b'PSA\000KEY\000'
 
diff --git a/tests/scripts/generate_psa_tests.py b/tests/scripts/generate_psa_tests.py
index 798b9ed..d8748db 100755
--- a/tests/scripts/generate_psa_tests.py
+++ b/tests/scripts/generate_psa_tests.py
@@ -526,7 +526,7 @@
 
     def keys_for_usage_extension(
             self,
-            implyer_usage: psa_storage.Expr,
+            implyer_usage: str,
             alg: str,
             key_type: str,
             params: Optional[Iterable[str]] = None
@@ -538,13 +538,13 @@
         keys = [] #type: List[StorageKey]
         kt = crypto_knowledge.KeyType(key_type, params)
         bits = kt.sizes_to_test()[0]
-        implicit = StorageKey.IMPLICIT_USAGE_FLAGS[implyer_usage]
+        implicit_usage = StorageKey.IMPLICIT_USAGE_FLAGS[implyer_usage]
         usage_flags = 'PSA_KEY_USAGE_EXPORT'
-        material_usage_flags = usage_flags + ' | ' + implyer_usage.string
-        expected_usage_flags = material_usage_flags + ' | ' + implicit.string
+        material_usage_flags = usage_flags + ' | ' + implyer_usage
+        expected_usage_flags = material_usage_flags + ' | ' + implicit_usage
         alg2 = 0
         key_material = kt.key_material(bits)
-        usage_expression = re.sub(r'PSA_KEY_USAGE_', r'', implyer_usage.string)
+        usage_expression = re.sub(r'PSA_KEY_USAGE_', r'', implyer_usage)
         alg_expression = re.sub(r'PSA_ALG_', r'', alg)
         alg_expression = re.sub(r',', r', ', re.sub(r' +', r'', alg_expression))
         key_type_expression = re.sub(r'\bPSA_(?:KEY_TYPE|ECC_FAMILY)_',
@@ -623,18 +623,13 @@
         alg_with_keys = self.gather_key_types_for_sign_alg()
         key_filter = StorageKey.IMPLICIT_USAGE_FLAGS_KEY_RESTRICTION
 
-        # Use a lookup for the extendable usage flags to able to sort them
-        usage_lookup = {} #type: Dict[str, psa_storage.Expr]
-        for usage_flag in StorageKey.IMPLICIT_USAGE_FLAGS:
-            usage_lookup[usage_flag.string] = usage_flag
-
         # Walk through all combintion. The key types must be filtered to fit
         # the specific usage flag.
         keys += [key
-                 for usage in sorted(usage_lookup)
+                 for usage in sorted(StorageKey.IMPLICIT_USAGE_FLAGS, key=str)
                  for alg in sorted(alg_with_keys)
                  for key_type in sorted(alg_with_keys[alg]) if re.match(key_filter[usage], key_type)
-                 for key in self.keys_for_usage_extension(usage_lookup[usage], alg, key_type)]
+                 for key in self.keys_for_usage_extension(usage, alg, key_type)]
 
         self.key_builder = prev_builder
         return keys