diff --git a/python/samba/tests/krb5/kdc_base_test.py b/python/samba/tests/krb5/kdc_base_test.py
index 16e3f7a6a73..7d180380d13 100644
--- a/python/samba/tests/krb5/kdc_base_test.py
+++ b/python/samba/tests/krb5/kdc_base_test.py
@@ -561,7 +561,7 @@ class KDCBaseTest(RawKerberosTest):
 
         return bind, identifier, attributes
 
-    def get_keys(self, samdb, dn):
+    def get_keys(self, samdb, dn, expected_etypes=None):
         admin_creds = self.get_admin_creds()
 
         bind, identifier, attributes = self.get_secrets(
@@ -599,9 +599,10 @@ class KDCBaseTest(RawKerberosTest):
                 pwd = attr.value_ctr.values[0].blob
                 keys[kcrypto.Enctype.RC4] = pwd.hex()
 
-        default_enctypes = self.get_default_enctypes()
+        if expected_etypes is None:
+            expected_etypes = self.get_default_enctypes()
 
-        self.assertCountEqual(default_enctypes, keys)
+        self.assertCountEqual(expected_etypes, keys)
 
         return keys
 
diff --git a/python/samba/tests/krb5/protected_users_tests.py b/python/samba/tests/krb5/protected_users_tests.py
index dfa6021453f..a03ccaf0c66 100755
--- a/python/samba/tests/krb5/protected_users_tests.py
+++ b/python/samba/tests/krb5/protected_users_tests.py
@@ -384,11 +384,10 @@ class ProtectedUsersTests(KDCBaseTest):
 
         client_creds.set_password(new_password)
 
-        keys = self.get_keys(samdb, client_dn)
-        self.assertEqual({kcrypto.Enctype.AES256,
-                          kcrypto.Enctype.AES128,
-                          kcrypto.Enctype.RC4},
-                         keys.keys())
+        self.get_keys(samdb, client_dn,
+                      expected_etypes={kcrypto.Enctype.AES256,
+                                       kcrypto.Enctype.AES128,
+                                       kcrypto.Enctype.RC4})
 
     # Test that DES-CBC-CRC cannot be used whether or not the user is
     # protected.