1 files changed, 4 insertions, 2 deletions
diff --git a/src/test/hs_ntor_ref.py b/src/test/hs_ntor_ref.py
index 1b9772a5d6..d58ac3ca23 100644
--- a/src/test/hs_ntor_ref.py
+++ b/src/test/hs_ntor_ref.py
@@ -65,14 +65,16 @@ except ImportError:
 try:
     # Pull the sha3 functions in.
     from hashlib import sha3_256, shake_256
-    shake_squeeze = shake_256.digest
+    def shake_squeeze(obj, n):
+        return obj.digest(n)
 except ImportError:
     if hasattr(sha3, "SHA3256"):
         # If this happens, then we have the old "sha3" module which
         # hashlib and pysha3 superseded.
         sha3_256 = sha3.SHA3256
         shake_256 = sha3.SHAKE256
-        shake_squeeze = shake_256.squeeze
+        def shake_squeeze(obj, n):
+            return obj.squeeze(n)
     else:
         # error code 77 tells automake to skip this test
         sys.exit(77)