1 files changed, 50 insertions, 1 deletions
diff --git a/tests/test_sftp_big.py b/tests/test_sftp_big.py
index acfe71e3..7d1110c3 100644
--- a/tests/test_sftp_big.py
+++ b/tests/test_sftp_big.py
@@ -30,7 +30,7 @@ import time
 
 from paramiko.common import o660
 
-from ._util import slow
+from ._util import slow, wait_until
 
 
 @slow
@@ -365,3 +365,52 @@ class TestBigSFTP:
         finally:
             sftp.remove(f"{sftp.FOLDER}/hongry.txt")
             t.packetizer.REKEY_BYTES = pow(2, 30)
+
+    def test_prefetch_limit(self, sftp):
+        """
+        write a 1MB file and prefetch with a limit
+        """
+        kblob = 1024 * b"x"
+        start = time.time()
+
+        def expect_prefetch_extents(file, expected_extents):
+            with file._prefetch_lock:
+                assert len(file._prefetch_extents) == expected_extents
+
+        try:
+            with sftp.open(f"{sftp.FOLDER}/hongry.txt", "w") as f:
+                for n in range(1024):
+                    f.write(kblob)
+                    if n % 128 == 0:
+                        sys.stderr.write(".")
+            sys.stderr.write(" ")
+
+            assert (
+                sftp.stat(f"{sftp.FOLDER}/hongry.txt").st_size == 1024 * 1024
+            )
+            end = time.time()
+            sys.stderr.write(f"{round(end - start)}s")
+
+            # read with prefetch, no limit
+            # expecting 32 requests (32k * 32 == 1M)
+            with sftp.open(f"{sftp.FOLDER}/hongry.txt", "rb") as f:
+                file_size = f.stat().st_size
+                f.prefetch(file_size)
+                wait_until(lambda: expect_prefetch_extents(f, 32))
+
+            # read with prefetch, limiting to 5 simultaneous requests
+            with sftp.open(f"{sftp.FOLDER}/hongry.txt", "rb") as f:
+                file_size = f.stat().st_size
+                f.prefetch(file_size, 5)
+                wait_until(lambda: expect_prefetch_extents(f, 5))
+                for n in range(1024):
+                    with f._prefetch_lock:
+                        assert len(f._prefetch_extents) <= 5
+                    data = f.read(1024)
+                    assert data == kblob
+
+                    if n % 128 == 0:
+                        sys.stderr.write(".")
+
+        finally:
+            sftp.remove(f"{sftp.FOLDER}/hongry.txt")