diff --git a/tests/scripts/generate_psa_wrappers.py b/tests/scripts/generate_psa_wrappers.py
index a1c775adc0..7fc852edc1 100755
--- a/tests/scripts/generate_psa_wrappers.py
+++ b/tests/scripts/generate_psa_wrappers.py
@@ -153,7 +153,9 @@ class PSAWrapperGenerator(c_wrapper_generator.Base):
         # Proof-of-concept: just instrument one function for now
         if function_name == 'psa_cipher_encrypt':
             return True
-        if function_name == 'psa_import_key' or function_name == 'psa_export_key' or function_name == 'psa_export_public_key':
+            if function_name in ('psa_import_key',
+                                 'psa_export_key',
+                                 'psa_export_public_key'):
             return True
         if function_name in ('psa_sign_message',
                              'psa_verify_message',