diff --git a/scripts/generate_psa_tests.py b/scripts/generate_psa_tests.py
index 417fc08f1..d18411662 100755
--- a/scripts/generate_psa_tests.py
+++ b/scripts/generate_psa_tests.py
@@ -41,6 +41,7 @@ def test_case_for_key_type_not_supported(
     tc.set_description('PSA {} {} {}-bit {} supported'
                        .format(verb, short_key_type, bits, adverb))
     tc.set_function(verb + '_not_supported')
+    tc.set_key_bits(bits)
     tc.set_arguments([key_type] + list(args))
     tc.set_dependencies(dependencies)
     tc.skip_if_any_not_implemented(dependencies)
@@ -153,6 +154,7 @@ def test_case_for_key_generation(
     tc.set_description('PSA {} {}-bit'
                        .format(short_key_type, bits))
     tc.set_function('generate_key')
+    tc.set_key_bits(bits)
     tc.set_arguments([key_type] + list(args) + [result])
     tc.set_dependencies(dependencies)
     tc.skip_if_any_not_implemented(dependencies)
@@ -279,7 +281,9 @@ class OpFail:
         tc.set_function(category.name.lower() + '_fail')
         arguments = [] # type: List[str]
         if kt:
-            key_material = kt.key_material(kt.sizes_to_test()[0])
+            bits = kt.sizes_to_test()[0]
+            tc.set_key_bits(bits)
+            key_material = kt.key_material(bits)
             arguments += [key_type, test_case.hex_string(key_material)]
         arguments.append(alg.expression)
         if category.is_asymmetric():
@@ -498,6 +502,7 @@ class StorageFormat:
         dependencies += psa_information.generate_deps_from_description(key.description)
         dependencies = psa_information.fix_key_pair_dependencies(dependencies, 'BASIC')
         tc.set_function('key_storage_' + verb)
+        tc.set_key_bits(key.bits)
         if self.forward:
             extra_arguments = []
         else:
diff --git a/scripts/mbedtls_framework/psa_test_case.py b/scripts/mbedtls_framework/psa_test_case.py
index 57172053e..d478d0ac3 100644
--- a/scripts/mbedtls_framework/psa_test_case.py
+++ b/scripts/mbedtls_framework/psa_test_case.py
@@ -70,11 +70,25 @@ class TestCase(test_case.TestCase):
         self.manual_dependencies = [] #type: List[str]
         self.automatic_dependencies = set() #type: Set[str]
         self.dependency_prefix = dependency_prefix #type: Optional[str]
+        self.key_bits = None #type: Optional[int]
+
+    def set_key_bits(self, key_bits: Optional[int]) -> None:
+        """Use the given key size for automatic dependency generation.
+
+        Call this function before set_arguments() if relevant.
+
+        This is only relevant for ECC and DH keys. For other key types,
+        this information is ignored.
+        """
+        self.key_bits = key_bits
 
     def infer_dependencies(self, arguments: List[str]) -> List[str]:
         """Infer dependencies based on the test case arguments."""
         dependencies = psa_information.automatic_dependencies(*arguments,
                                                               prefix=self.dependency_prefix)
+        if self.key_bits is not None:
+            dependencies = psa_information.finish_family_dependencies(dependencies,
+                                                                      self.key_bits)
         return dependencies
 
     def set_arguments(self, arguments: List[str]) -> None: