diff options
| -rw-r--r-- | tests/_util.py | 23 | ||||
| -rw-r--r-- | tests/test_sftp_big.py | 14 |
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 |