diff --git a/tests/scripts/generate_psa_wrappers.py b/tests/scripts/generate_psa_wrappers.py
index 0b7cf44..92c9310 100755
--- a/tests/scripts/generate_psa_wrappers.py
+++ b/tests/scripts/generate_psa_wrappers.py
@@ -171,6 +171,9 @@
                              'psa_hash_compute',
                              'psa_hash_compare'):
             return True
+        if function_name in ('psa_key_derivation_key_agreement',
+                             'psa_raw_key_agreement'):
+            return True
         return False
 
     def _write_function_call(self, out: typing_util.Writable,