Make test for prefetch requests limit stable

The prefetch requests are added in a background process, so the might
not be ready when the test wants to count them. Fix it by introducing a
wait_for untility function that waits for a assert to pass for a
specified timeout, to give the background thread the chance to start and
add the prefetch requests.

Also lock the prefetch lock before trying to access the extents to
prevent a data race.

2 files changed, 31 insertions, 6 deletions

diff --git a/tests/_util.py b/tests/_util.py
index eaf6aac4..ec7585df 100644
--- a/tests/_util.py
+++ b/tests/_util.py
@@ -7,7 +7,7 @@ import socket
 import struct
 import sys
 import unittest
-from time import sleep
+import time
 import threading
 
 import pytest
@@ -303,7 +303,7 @@ class TestServer(ServerInterface):
         if username == "bad-server":
             raise Exception("Ack!")
         if username == "unresponsive-server":
-            sleep(5)
+            time.sleep(5)
             return AUTH_SUCCESSFUL
         return AUTH_FAILED
 
@@ -439,3 +439,22 @@ def server(
     ts.close()
     socks.close()
     sockc.close()
+
+
+def wait_until(condition, *, timeout=2):
+    """
+    Wait until `condition()` no longer raises an `AssertionError` or until
+    `timeout` seconds have passed, which causes a `TimeoutError` to be raised.
+    """
+    deadline = time.time() + timeout
+
+    while True:
+        try:
+            condition()
+        except AssertionError as e:
+            if time.time() > deadline:
+                timeout_message = f"Condition not reached after {timeout}s"
+                raise TimeoutError(timeout_message) from e
+        else:
+            return
+        time.sleep(0.01)
diff --git a/tests/test_sftp_big.py b/tests/test_sftp_big.py
index d3e0cb1b..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
@@ -372,6 +372,11 @@ class TestBigSFTP:
         """
         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):
@@ -391,15 +396,16 @@ class TestBigSFTP:
             with sftp.open(f"{sftp.FOLDER}/hongry.txt", "rb") as f:
                 file_size = f.stat().st_size
                 f.prefetch(file_size)
-                assert len(f._prefetch_extents) == 32
+                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)
-                assert len(f._prefetch_extents) == 5
+                wait_until(lambda: expect_prefetch_extents(f, 5))
                 for n in range(1024):
-                    assert len(f._prefetch_extents) <= 5
+                    with f._prefetch_lock:
+                        assert len(f._prefetch_extents) <= 5
                     data = f.read(1024)
                     assert data == kblob