From 526f91cb2a9c7dd60dcfd50e4132c3260aab7871 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:34:56 +0900 Subject: [PATCH 01/19] Remove and merge S3 error, path and pairing tests that other tests cover (GH-1239) - Fold the unknown error code into the errno table, which now asserts that an error without a status has no http_status_code, and remove the test of the S3ClientError properties. - Remove the version_id spelling test of S3Path, which the parse and name tests imply. - Fold the ordinary move of test_move_pairs into the null-version bucket-state test. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_errors.py | 21 +++++-------------- tests/pyathena/filesystem/test_s3_path.py | 3 --- .../filesystem/test_s3_path_pairing.py | 18 ++++++++-------- 3 files changed, 14 insertions(+), 28 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_errors.py b/tests/pyathena/filesystem/test_s3_errors.py index 246f2cb30..2554d442a 100644 --- a/tests/pyathena/filesystem/test_s3_errors.py +++ b/tests/pyathena/filesystem/test_s3_errors.py @@ -23,15 +23,6 @@ def _client_error(code, message="error message", status_code=None, condition=Non class TestS3ClientError: - def test_properties(self): - error = S3ClientError(_client_error("NoSuchKey", status_code=404)) - assert error.code == "NoSuchKey" - assert error.message == "error message" - assert error.http_status_code == 404 - - error = S3ClientError(_client_error("NoSuchKey")) - assert error.http_status_code is None - @pytest.mark.parametrize( ("code", "expected"), [ @@ -63,10 +54,14 @@ def test_os_error_by_error_code(self, code, expected): ("ServiceUnavailable", errno.EBUSY), ("OperationAborted", errno.EBUSY), ("InternalError", errno.EIO), + # An unknown code without a status code. + ("UnknownCode", errno.EIO), ], ) def test_os_error_by_error_code_errno(self, code, expected_errno): - actual = S3ClientError(_client_error(code)).os_error + error = S3ClientError(_client_error(code)) + assert error.http_status_code is None + actual = error.os_error # OSError specializes itself into a subclass for some errno values # (e.g., EPERM -> PermissionError), so check the errno, not the exact type. assert isinstance(actual, OSError) @@ -91,12 +86,6 @@ def test_os_error_by_http_status_code(self, status_code, expected_errno): assert isinstance(actual, OSError) assert actual.errno == expected_errno - def test_os_error_unknown_code(self): - actual = S3ClientError(_client_error("UnknownCode")).os_error - assert type(actual) is OSError - assert actual.errno == errno.EIO - assert "error message" in str(actual) - @pytest.mark.parametrize( ("condition", "expected"), [ diff --git a/tests/pyathena/filesystem/test_s3_path.py b/tests/pyathena/filesystem/test_s3_path.py index 6f3a13db0..efa4b906b 100644 --- a/tests/pyathena/filesystem/test_s3_path.py +++ b/tests/pyathena/filesystem/test_s3_path.py @@ -115,9 +115,6 @@ def test_names(self, path, name, string, uri): assert str(path) == string assert path.uri == uri - def test_str_spells_the_version_query_as_version_id(self): - assert str(S3Path.parse("s3a://bucket/key?version_id=v1")) == "bucket/key?versionId=v1" - @pytest.mark.parametrize( ("path", "expected"), [ diff --git a/tests/pyathena/filesystem/test_s3_path_pairing.py b/tests/pyathena/filesystem/test_s3_path_pairing.py index 4b549352b..6685a4dc2 100644 --- a/tests/pyathena/filesystem/test_s3_path_pairing.py +++ b/tests/pyathena/filesystem/test_s3_path_pairing.py @@ -134,32 +134,32 @@ def test_copy_pairs_needs_sources(self): with pytest.raises(ValueError, match="sources"): S3PathPairing("s3://bucket/a/", "s3://bucket/b/").copy_pairs() - def test_move_pairs(self): - # The "null" version of a key moved onto the key stays in place. - pairs = [("s3://bucket/b?versionId=null", "s3://bucket/b"), ("bucket/d/a", "s3://bucket/z")] - - assert _pairing(pairs).conflict_candidates(pairs) == [] - assert _pairing(pairs).move_pairs(pairs) == [("bucket/d/a", "s3://bucket/z")] - @pytest.mark.parametrize( ("enabled_buckets", "expected"), [ + # The "null" version of a key moved onto the key stays in place. ((), []), (("enabled",), [("enabled/key?versionId=null", "enabled/key")]), (("other",), [("other/key?versionId=null", "other/key")]), ], ) def test_move_pairs_null_version_bucket_state(self, enabled_buckets, expected): + moved = ("enabled/d/a", "s3://enabled/z") pairs = [ ("enabled/key?versionId=null", "enabled/key"), ("other/key?versionId=null", "other/key"), ("enabled/unchanged?versionId=null", "s3://enabled/unchanged?versionID=null"), + # An ordinary move. + moved, ] pairing = _pairing(pairs) assert pairing.conflict_candidates(pairs, versioning_enabled_buckets=enabled_buckets) == [] - assert pairing.move_pairs(pairs, versioning_enabled_buckets=enabled_buckets) == expected - assert pairing.move_pairs(pairs) == [] + assert pairing.move_pairs(pairs, versioning_enabled_buckets=enabled_buckets) == [ + *expected, + moved, + ] + assert pairing.move_pairs(pairs) == [moved] def test_move_pairs_distinguishes_null_version_source(self): pairs = [ From eb19f3fcf5dda2b430de082fc0e7d8b4d7f48374 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:38:24 +0900 Subject: [PATCH 02/19] Remove and merge S3Core tests that other core tests cover (GH-1239) - Narrow test_properties to the documented copy of request_kwargs. - Remove the per-operation request_kwargs and error translation tests, whose rule every operation shares through call(). - Fold the S3Metadata mapping of test_replace_object_metadata_omits_ unset_fields into test_replace_object_metadata; setxattr() checks the omitted fields. - Fold the keyless path into TestS3DeleteBatch::test_from_paths and the storage class into test_list_objects, and remove the frozen dataclass test. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_core.py | 121 ++++------------------ 1 file changed, 23 insertions(+), 98 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index 2a0695b30..2d2915b3e 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -83,28 +83,11 @@ def test_operation_params(self, method, params, expected): core, _ = _make_core() assert core.operation_params(method, params) == expected - def test_properties(self): + def test_request_kwargs_is_a_copy(self): core, _ = _make_core(request_kwargs={"RequestPayer": "requester"}) - assert core.retry_config.attempt == 1 - assert core.request_kwargs == {"RequestPayer": "requester"} - # A copy, so the core keeps its parameters. + # A copy, so a caller that changes it leaves the core's parameters. core.request_kwargs.clear() assert core.request_kwargs == {"RequestPayer": "requester"} - assert S3Core(core.client).retry_config.attempt == RetryConfig().attempt - - def test_request_kwargs_go_to_operations_that_accept_them(self): - core, stubber = _make_core(request_kwargs={"RequestPayer": "requester"}) - stubber.add_response( - "head_object", - {"ContentLength": 1}, - {"Bucket": "bucket", "Key": "key", "RequestPayer": "requester"}, - ) - # HeadBucket does not accept RequestPayer. - stubber.add_response("head_bucket", {}, {"Bucket": "bucket"}) - with stubber: - core.head_object(S3Path("bucket", "key")) - core.head_bucket("bucket") - stubber.assert_no_pending_responses() def test_call_params_are_sent_as_given(self): # A per-call parameter is not filtered, so botocore rejects a @@ -516,20 +499,20 @@ def test_get_bucket_versioning(self, response, expected): stubber.assert_no_pending_responses() assert status == expected - def test_get_bucket_versioning_translates_errors(self): - core, stubber = _make_core() - stubber.add_client_error( - "get_bucket_versioning", service_error_code="NoSuchBucket", http_status_code=404 - ) - with stubber, pytest.raises(FileNotFoundError): - core.get_bucket_versioning("bucket") - def test_list_objects(self): core, stubber = _make_core() stubber.add_response( "list_objects_v2", { - "Contents": [{"Key": "dir/a", "Size": 1, "ETag": '"a"', "LastModified": MODIFIED}], + "Contents": [ + { + "Key": "dir/a", + "Size": 1, + "ETag": '"a"', + "LastModified": MODIFIED, + "StorageClass": "STANDARD_IA", + } + ], "CommonPrefixes": [{"Prefix": "dir/sub/"}], "KeyCount": 2, "IsTruncated": True, @@ -556,7 +539,12 @@ def test_list_objects(self): bucket="bucket", objects=( S3ObjectSummary( - bucket="bucket", key="dir/a", size=1, etag='"a"', last_modified=MODIFIED + bucket="bucket", + key="dir/a", + size=1, + etag='"a"', + last_modified=MODIFIED, + storage_class="STANDARD_IA", ), ), common_prefixes=(S3CommonPrefix("bucket", "dir/sub/"),), @@ -873,14 +861,6 @@ def test_delete_objects_quiet(self): result = core.delete_objects(S3DeleteBatch("bucket", (S3Path("bucket", "a"),))) assert result == S3DeleteResult("bucket") - def test_delete_objects_translates_errors(self): - core, stubber = _make_core() - stubber.add_client_error( - "delete_objects", service_error_code="AccessDenied", http_status_code=403 - ) - with stubber, pytest.raises(PermissionError): - core.delete_objects(S3DeleteBatch("bucket", (S3Path("bucket", "a"),))) - def test_create_multipart_upload(self): core, stubber = _make_core() stubber.add_response( @@ -1776,46 +1756,14 @@ def test_replace_object_metadata(self): core.replace_object_metadata( S3Path("bucket", "key"), head, - {"new": "2"}, + # The user-defined metadata can be any mapping, such as an + # S3Metadata. + S3Metadata({"Metadata": {"new": "2"}}), ContentType="text/plain", SSECustomerAlgorithm="AES256", ) stubber.assert_no_pending_responses() - def test_replace_object_metadata_omits_unset_fields(self): - core, stubber = _make_core() - stubber.add_response( - "copy_object", - {}, - { - "CopySource": {"Bucket": "bucket", "Key": "key"}, - "Bucket": "bucket", - "Key": "key", - "Metadata": {}, - "MetadataDirective": "REPLACE", - # S3Metadata reports STANDARD when HeadObject omits it. - "StorageClass": "STANDARD", - }, - ) - # The user-defined metadata can be any mapping, such as an S3Metadata. - stubber.add_response( - "copy_object", - {}, - { - "CopySource": {"Bucket": "bucket", "Key": "key"}, - "Bucket": "bucket", - "Key": "key", - "Metadata": {"a": "1"}, - "MetadataDirective": "REPLACE", - "StorageClass": "STANDARD", - }, - ) - with stubber: - core.replace_object_metadata(S3Path("bucket", "key"), S3Metadata({}), {}) - head = S3Metadata({"Metadata": {"a": "1"}}) - core.replace_object_metadata(S3Path("bucket", "key"), head, head) - stubber.assert_no_pending_responses() - @pytest.mark.parametrize( "name", ["Metadata", "MetadataDirective", "CopySource", "Bucket", "Key"] ) @@ -1890,6 +1838,9 @@ def test_from_paths(self): assert batches[1].objects == (S3Path("b1", f"k{S3DeleteBatch.MAX_KEYS}"),) assert batches[2].objects == (S3Path("b2", "a", "v1"),) assert S3DeleteBatch.from_paths([]) == [] + # A path without a key is rejected, also among other paths. + with pytest.raises(ValueError, match="Not an object of the bucket b1: s3://b1"): + S3DeleteBatch.from_paths([S3Path("b1", "a"), S3Path("b1")]) @pytest.mark.parametrize( ("objects", "match"), @@ -1908,10 +1859,6 @@ def test_invalid(self, objects, match): with pytest.raises(ValueError, match=match): S3DeleteBatch("bucket", objects) - def test_from_paths_rejects_bucket_paths(self): - with pytest.raises(ValueError, match="Not an object of the bucket bucket: s3://bucket"): - S3DeleteBatch.from_paths([S3Path("bucket", "a"), S3Path("bucket")]) - class TestS3DeleteError: def test_from_response(self): @@ -1921,25 +1868,3 @@ def test_from_response(self): assert error == S3DeleteError(S3Path("bucket", "a", "v1"), "InternalError", "Error") assert str(error) == "bucket/a?versionId=v1 (InternalError: Error)" assert str(S3DeleteError.from_response("bucket", {"Key": "a"})) == "bucket/a (None: None)" - - -class TestS3ObjectSummary: - def test_from_response(self): - entry = { - "Key": "dir/a", - "Size": 1, - "ETag": '"a"', - "LastModified": MODIFIED, - "StorageClass": "STANDARD_IA", - } - summary = S3ObjectSummary.from_response("bucket", entry) - assert summary == S3ObjectSummary("bucket", "dir/a", 1, '"a"', MODIFIED, "STANDARD_IA") - assert summary.path == S3Path("bucket", "dir/a") - # Optional fields that the entry does not have are None. - assert S3ObjectSummary.from_response("bucket", {"Key": "k"}) == S3ObjectSummary( - "bucket", "k" - ) - - def test_frozen(self): - with pytest.raises(AttributeError): - S3ObjectSummary("bucket", "k").key = "other" # type: ignore[misc] From b7aa10f024bff14341f77c75d65e014c51f3e256 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:43:46 +0900 Subject: [PATCH 03/19] Merge S3Core path, ACL, multipart and presigned URL tests (GH-1239) - Fold the per-method path checks into test_object_operations_reject_ paths and the create_bucket ACLs into test_acl_validation. - Keep one missing identity field per multipart method beyond the three of upload_part. - Fold test_upload_part into the creation-algorithm test, the ClientError cause and InvalidRange of get_object into test_call_translates_errors, and the precedence of the given parts into test_complete_multipart_ upload; the per-field checksum mapping is S3MultipartUploadPart's. - Sign presigned URLs locally and assert the URL instead of the arguments of a mocked call. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_core.py | 281 ++++++---------------- 1 file changed, 78 insertions(+), 203 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index 2d2915b3e..ef33479c1 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -9,12 +9,12 @@ import gc import io import re +import time import weakref from concurrent.futures import Future from datetime import UTC, datetime from itertools import pairwise from typing import Any -from unittest import mock from urllib.parse import parse_qs, urlsplit import boto3 @@ -103,13 +103,18 @@ def test_call_params_are_sent_as_given(self): (404, "NoSuchKey", FileNotFoundError), (403, "403", PermissionError), (500, "InternalError", OSError), + (416, "InvalidRange", OSError), ], ) def test_call_translates_errors(self, status, code, error): core, stubber = _make_core() stubber.add_client_error("head_object", service_error_code=code, http_status_code=status) - with stubber, pytest.raises(error): + with stubber, pytest.raises(error) as e: core.head_object(S3Path("bucket", "key")) + # The ClientError is the cause, whose code callers such as reads of + # an InvalidRange read. + assert isinstance(e.value.__cause__, botocore.exceptions.ClientError) + assert e.value.__cause__.response["Error"]["Code"] == code @pytest.mark.parametrize("translated", [True, False]) def test_call_failure_frees_request(self, translated): @@ -205,11 +210,6 @@ def test_head_object(self): metadata.object_lock_legal_hold_status, ) == ("GOVERNANCE", MODIFIED, "ON") - def test_head_object_requires_a_key(self): - core, _ = _make_core() - with pytest.raises(ValueError, match="no key"): - core.head_object(S3Path("bucket")) - def test_head_bucket(self): core, stubber = _make_core() stubber.add_response("head_bucket", {"BucketRegion": "us-west-2"}, {"Bucket": "bucket"}) @@ -292,18 +292,6 @@ def test_get_object_closes_body(self): stubber.assert_no_pending_responses() assert [raw.closed for raw in raws] == [True, True] - @pytest.mark.parametrize( - ("status", "code", "error"), - [(404, "NoSuchKey", FileNotFoundError), (416, "InvalidRange", OSError)], - ) - def test_get_object_translates_errors(self, status, code, error): - core, stubber = _make_core() - stubber.add_client_error("get_object", service_error_code=code, http_status_code=status) - with stubber, pytest.raises(error) as e: - core.get_object(S3Path("bucket", "key"), (10, None)) - assert isinstance(e.value.__cause__, botocore.exceptions.ClientError) - assert e.value.__cause__.response["Error"]["Code"] == code - @pytest.mark.parametrize( ("body", "params", "expected"), [ @@ -330,19 +318,6 @@ def test_put_object(self, body, params, expected): stubber.assert_no_pending_responses() assert (result.etag, result.version_id) == ('"e"', "v1") - @pytest.mark.parametrize( - ("path", "match"), - [ - (S3Path("bucket"), "no key"), - (S3Path("bucket", "key", "v1"), "Cannot write to a version"), - (S3Path("bucket", "key", "null"), "Cannot write to a version"), - ], - ) - def test_put_object_rejects_paths(self, path, match): - core, stubber = _make_core() - with stubber, pytest.raises(ValueError, match=match): - core.put_object(path, b"data") - @pytest.mark.parametrize( ("client_region", "kwargs", "expected"), [ @@ -456,13 +431,6 @@ def test_create_bucket_configuration(self, client_region, region_name, configura # The given configuration is not modified. assert given == configuration - @pytest.mark.parametrize("acl", ["invalid", "bucket-owner-full-control"]) - def test_create_bucket_rejects_acls(self, acl): - # An object-only canned ACL is not accepted for a bucket either. - core, stubber = _make_core() - with stubber, pytest.raises(ValueError, match="ACL not in"): - core.create_bucket("bucket", acl=acl) - def test_delete_bucket(self): core, stubber = _make_core() stubber.add_response( @@ -814,11 +782,6 @@ def test_delete_object(self): core.delete_object(S3Path("bucket", "key", "v1"), BypassGovernanceRetention=True) stubber.assert_no_pending_responses() - def test_delete_object_requires_a_key(self): - core, _ = _make_core() - with pytest.raises(ValueError, match="has no key"): - core.delete_object(S3Path("bucket")) - def test_delete_objects(self): core, stubber = _make_core() stubber.add_response( @@ -875,30 +838,6 @@ def test_create_multipart_upload(self): ) assert upload.upload_id == "u" - def test_upload_part(self): - core, stubber = _make_core() - stubber.add_response( - "upload_part", - {"ETag": '"e1"'}, - { - "Bucket": "bucket", - "Key": "key", - "UploadId": "u", - "PartNumber": 1, - "Body": b"data", - "SSECustomerAlgorithm": "AES256", - }, - ) - with stubber: - part = core.upload_part( - S3MultipartUpload({"Bucket": "bucket", "Key": "key", "UploadId": "u"}), - 1, - b"data", - SSECustomerAlgorithm="AES256", - PartNumber=99, - ) - assert (part.part_number, part.etag) == (1, '"e1"') - @pytest.mark.parametrize( "bucket", ["myap-abc123-s3alias", "arn:aws:s3:us-east-1:123456789012:accesspoint/myap"] ) @@ -1017,6 +956,8 @@ def test_complete_multipart_upload(self): S3MultipartUpload({"Bucket": "bucket", "Key": "key", "UploadId": "u"}), parts, RequestPayer="requester", + # The given parts take precedence. + MultipartUpload={"Parts": []}, ) assert (completed.etag, completed.version_id) == ('"dst"', "v-dst") @@ -1042,37 +983,25 @@ def test_abort_multipart_upload(self): ) @pytest.mark.parametrize( - ("method", "args"), + ("method", "args", "missing", "message"), [ - ("create_multipart_upload", (S3Path("bucket"),)), - ( - "upload_part_copy", - ( - S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "u"}), - 1, - S3Path("bucket"), - ), - ), + ("upload_part", (1, b"data"), "Bucket", "no bucket"), + ("upload_part", (1, b"data"), "Key", "no key"), + ("upload_part", (1, b"data"), "UploadId", "no upload ID"), + # The other operations check the same identity. + ("upload_part_copy", (1, S3Path("source", "key")), "UploadId", "no upload ID"), + ("complete_multipart_upload", ([],), "Key", "no key"), + ("abort_multipart_upload", (), "Bucket", "no bucket"), ], - ) - def test_multipart_upload_requires_keys(self, method, args): - core, _ = _make_core() - with pytest.raises(ValueError, match="has no key"): - getattr(core, method)(*args) - - @pytest.mark.parametrize( - ("method", "args"), - [ - ("upload_part", (1, b"data")), - ("upload_part_copy", (1, S3Path("source", "key"))), - ("complete_multipart_upload", ([],)), - ("abort_multipart_upload", ()), + ids=[ + "upload_part-Bucket", + "upload_part-Key", + "upload_part-UploadId", + "upload_part_copy-UploadId", + "complete_multipart_upload-Key", + "abort_multipart_upload-Bucket", ], ) - @pytest.mark.parametrize( - ("missing", "message"), - [("Bucket", "no bucket"), ("Key", "no key"), ("UploadId", "no upload ID")], - ) def test_multipart_upload_requires_identity(self, method, args, missing, message): core, _ = _make_core() response = {"Bucket": "bucket", "Key": "key", "UploadId": "u"} @@ -1095,9 +1024,12 @@ def test_upload_part_uses_creation_algorithm_and_identity(self): "PartNumber": 1, "Body": b"data", "ChecksumAlgorithm": "SHA256", + "SSECustomerAlgorithm": "AES256", }, ) with stubber: + # The identity and the part number override the per-call fields, + # and the other parameters are sent. part = core.upload_part( upload, 1, @@ -1105,17 +1037,14 @@ def test_upload_part_uses_creation_algorithm_and_identity(self): Bucket="other", Key="other", UploadId="other", + PartNumber=99, ChecksumAlgorithm="CRC32", + SSECustomerAlgorithm="AES256", ) stubber.assert_no_pending_responses() + assert (part.part_number, part.etag) == (1, '"part"') assert part.to_api_repr()["ChecksumSHA256"] == "sha" - def test_create_multipart_upload_rejects_versions(self): - # A write replaces the object at the key, not the named version. - core, _ = _make_core() - with pytest.raises(ValueError, match="Cannot write to a version"): - core.create_multipart_upload(S3Path("bucket", "key", "v1")) - @pytest.mark.parametrize( ("size", "block_size", "ranges"), [ @@ -1478,76 +1407,6 @@ def test_copy_object_annotation(self): ) stubber.assert_no_pending_responses() - @pytest.mark.parametrize( - "field", - [ - "ChecksumCRC32", - "ChecksumCRC32C", - "ChecksumCRC64NVME", - "ChecksumSHA1", - "ChecksumSHA256", - "ChecksumSHA512", - "ChecksumMD5", - "ChecksumXXHASH64", - "ChecksumXXHASH3", - "ChecksumXXHASH128", - ], - ) - @pytest.mark.parametrize("copy", [False, True]) - def test_complete_multipart_upload_preserves_part_checksums(self, field, copy): - core, stubber = _make_core() - identity = {"Bucket": "bucket", "Key": "key", "UploadId": "u"} - algorithm = field.removeprefix("Checksum") - stubber.add_response( - "create_multipart_upload", - {**identity, "ChecksumAlgorithm": algorithm}, - {"Bucket": "bucket", "Key": "key", "ChecksumAlgorithm": algorithm}, - ) - expected_parts = [] - for number in (1, 2): - result = {"ETag": f'"e{number}"', field: f"checksum{number}"} - request = {"Bucket": "bucket", "Key": "key", "UploadId": "u", "PartNumber": number} - if copy: - request["CopySource"] = {"Bucket": "bucket", "Key": "source"} - stubber.add_response("upload_part_copy", {"CopyPartResult": result}, request) - else: - request["Body"] = b"data" - request["ChecksumAlgorithm"] = algorithm - request[field] = f"checksum{number}" - stubber.add_response("upload_part", result, request) - expected_parts.append({**result, "PartNumber": number}) - stubber.add_response( - "complete_multipart_upload", - {"ETag": '"done"'}, - { - "Bucket": "bucket", - "Key": "key", - "UploadId": "u", - "MultipartUpload": {"Parts": expected_parts}, - "RequestPayer": "requester", - }, - ) - with stubber: - upload = core.create_multipart_upload( - S3Path("bucket", "key"), ChecksumAlgorithm=algorithm - ) - if copy: - parts = [ - core.upload_part_copy(upload, n, S3Path("bucket", "source")) for n in (1, 2) - ] - else: - parts = [ - core.upload_part(upload, n, b"data", **{field: f"checksum{n}"}) for n in (1, 2) - ] - completed = core.complete_multipart_upload( - upload, - parts, - RequestPayer="requester", - MultipartUpload={"Parts": []}, - ) - stubber.assert_no_pending_responses() - assert completed.etag == '"done"' - @pytest.mark.parametrize( ("algorithm", "checksum_type"), [(None, None), ("SHA256", "COMPOSITE"), ("CRC32", "FULL_OBJECT")], @@ -1665,6 +1524,8 @@ def test_put_bucket_acl(self): ("put_object_acl", (S3Path("bucket", "key"), "invalid")), # An object ACL that a bucket does not accept. ("put_bucket_acl", ("bucket", "bucket-owner-full-control")), + ("create_bucket", ("bucket", "invalid")), + ("create_bucket", ("bucket", "bucket-owner-full-control")), ], ) def test_acl_validation(self, method, args): @@ -1675,6 +1536,31 @@ def test_acl_validation(self, method, args): @pytest.mark.parametrize( ("method", "args", "match"), [ + ("head_object", (S3Path("bucket"),), "has no key"), + ("delete_object", (S3Path("bucket"),), "has no key"), + ("put_object", (S3Path("bucket"), b"data"), "has no key"), + # A write replaces the object at the key, not the named version. + ("put_object", (S3Path("bucket", "key", "v1"), b"data"), "Cannot write to a version"), + ( + "put_object", + (S3Path("bucket", "key", "null"), b"data"), + "Cannot write to a version", + ), + ("create_multipart_upload", (S3Path("bucket"),), "has no key"), + ( + "create_multipart_upload", + (S3Path("bucket", "key", "v1"),), + "Cannot write to a version", + ), + ( + "upload_part_copy", + ( + S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "u"}), + 1, + S3Path("bucket"), + ), + "has no key", + ), ("get_object_tagging", (S3Path("bucket"),), "has no key"), ("put_object_tagging", (S3Path("bucket"), {}), "has no key"), ("put_object_acl", (S3Path("bucket"), "private"), "has no key"), @@ -1773,53 +1659,42 @@ def test_replace_object_metadata_rejects_fields_of_the_copy(self, name): core.replace_object_metadata(S3Path("bucket", "key"), S3Metadata({}), {}, **{name: "x"}) @pytest.mark.parametrize( - ("path", "client_method", "params", "expected"), + ("path", "client_method", "params", "key", "query"), [ - (S3Path("bucket", "key"), "get_object", {}, {"Bucket": "bucket", "Key": "key"}), + (S3Path("bucket", "key"), "get_object", {}, "key", {}), ( S3Path("bucket", "key", "null"), "get_object", {"ResponseContentType": "text/csv"}, - { - "Bucket": "bucket", - "Key": "key", - "VersionId": "null", - "ResponseContentType": "text/csv", - }, + "key", + {"versionId": ["null"], "response-content-type": ["text/csv"]}, ), # The parameters take precedence over the path. ( S3Path("bucket", "key", "v1"), - "put_object", - {"Key": "other", "VersionId": None}, - {"Bucket": "bucket", "Key": "other", "VersionId": None}, + "get_object", + {"Key": "other"}, + "other", + {"versionId": ["v1"]}, ), + (S3Path("bucket", "key"), "put_object", {"Key": "other"}, "other", {}), ], ) - def test_generate_presigned_url(self, path, client_method, params, expected): - core, _ = _make_core(request_kwargs={"RequestPayer": "requester"}) - core.call = mock.MagicMock(return_value="https://signed") - assert ( - core.generate_presigned_url(path, client_method, expires_in=60, **params) - == "https://signed" - ) - # request_kwargs are not signed. - core.call.assert_called_once_with( - core.client.generate_presigned_url, - ClientMethod=client_method, - Params=expected, - ExpiresIn=60, - ) - - def test_generate_presigned_url_signs_locally(self): + def test_generate_presigned_url(self, path, client_method, params, key, query): # No request is sent, so the stubber has no responses to return. core, stubber = _make_core(request_kwargs={"RequestPayer": "requester"}) with stubber: - url = core.generate_presigned_url(S3Path("bucket", "key", "v1")) - query = parse_qs(urlsplit(url).query) - assert query["versionId"] == ["v1"] - # request_kwargs are not signed. - assert not any(k.lower() == "x-amz-request-payer" for k in query) + url = core.generate_presigned_url(path, client_method, expires_in=3600, **params) + signed = time.time() + parts = urlsplit(url) + assert (parts.netloc, parts.path) == ("bucket.s3.amazonaws.com", f"/{key}") + signed_query = parse_qs(parts.query) + # The signature fields of the client's signer, and the expiry. + expires = int(signed_query.pop("Expires")[0]) + assert signed + 3600 - 60 < expires <= signed + 3600 + del signed_query["AWSAccessKeyId"], signed_query["Signature"] + # request_kwargs, such as RequestPayer, are not signed. + assert signed_query == query class TestS3DeleteBatch: From 9ce11d0fef4a48469403209d0df2730ad306ff3a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:54:47 +0900 Subject: [PATCH 04/19] Remove and rewrite the offline async filesystem tests (GH-1239) Apply the offline TestAioS3FileSystem, TestAioS3File and TestAioS3StreamedFile rows of #1241: remove the tests whose contract the shared sync, core and executor tests check, check the async single-request copy, recursive rm with maxdepth and glob listing at the request level, drive the transaction put_file() through the public call, and cancel a real read instead of patching asyncio.wait. Keep test_put_file_mode, the non-transaction cases of test_pipe_file_compression and test_commit_failed_part_first, and merge test_open_version_id into test_open_lookup_parameters: each is the only test that fails when its async code is broken. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 759 ++++++--------------- 1 file changed, 198 insertions(+), 561 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 7868fc201..5beb2e9f3 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -24,7 +24,7 @@ from botocore.stub import Stubber from fsspec import Callback -from pyathena.filesystem.s3 import S3File, S3FileSystem +from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile from pyathena.filesystem.s3_core import S3Core, S3MultipartCopyPlan from pyathena.filesystem.s3_executor import S3AioThreadPool, S3Executor @@ -37,7 +37,7 @@ ) from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_writer import S3MultipartStateError -from pyathena.util import RetryConfig, stop_retries_on +from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect from tests.pyathena.util import ( @@ -63,49 +63,6 @@ def register_async_filesystem(): @pytest.mark.usefixtures("register_async_filesystem") class TestAioS3FileSystem: - @pytest.mark.parametrize("max_workers", [1, 4]) - @pytest.mark.asyncio - async def test_copy_object_with_multipart_upload_part_sizes(self, max_workers): - # GH-951: the parts are within the S3 part size limits whatever the - # number of workers; a single worker used to copy the whole object - # as one part larger than 5 GiB. - fs = AioS3FileSystem( - connection=mock.MagicMock(), max_workers=max_workers, skip_instance_cache=True - ) - sync_fs = fs._sync_fs - sync_fs.core.create_multipart_upload = mock.MagicMock( - return_value=S3MultipartUpload( - {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} - ) - ) - sync_fs.core.upload_part_copy = mock.MagicMock( - side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) - ) - sync_fs.core.complete_multipart_upload = mock.MagicMock() - # The HeadObject of the source. - sync_fs._call = sync_fs._core.call = mock.MagicMock( - return_value={"ContentLength": 5 * 2**30 + 2**20} - ) - - await fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), - S3Path("bucket", "dst"), - # Copy without reading the metadata, tags and annotations of the - # source (GH-973). - MetadataDirective="REPLACE", - TaggingDirective="REPLACE", - AnnotationDirective="EXCLUDE", - ) - - parts = sorted( - (c.kwargs["part_number"], c.kwargs["range_"]) - for c in sync_fs.core.upload_part_copy.call_args_list - ) - assert parts == [ - (1, (0, 5 * 2**29 + 2**19)), - (2, (5 * 2**29 + 2**19, 5 * 2**30 + 2**20)), - ] - @staticmethod async def _multipart_copy(fs=None, **kwargs): # max_workers=1 runs the stubbed requests in a deterministic order. @@ -133,30 +90,31 @@ async def _multipart_copy(fs=None, **kwargs): async def test_copy_object_with_multipart_upload_small_head_object_size(self, size): # GH-973: see # TestS3FileSystem.test_copy_object_with_multipart_upload_small_head_object_size. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - sync_fs = fs._sync_fs - sync_fs._call = sync_fs._core.call = mock.MagicMock( - return_value={"ContentLength": size, "VersionId": "v1"} - ) - sync_fs.core.copy_object = mock.MagicMock() - sync_fs.core.create_multipart_upload = mock.MagicMock() - - await fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), - S3Path("bucket", "dst"), - ContentType="text/csv", - RequestPayer="requester", - ) - - sync_fs.core.copy_object.assert_called_once_with( - S3Path("bucket", "src", "v1"), - S3Path("bucket", "dst"), - ContentType="text/csv", - RequestPayer="requester", + # The async copy has its own single-request branch: the reported + # version is copied with one CopyObject and nothing else is sent. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True ) - sync_fs.core.create_multipart_upload.assert_not_called() - # Only HeadObject; the tags are not read for the multipart upload. - assert sync_fs._call.call_count == 1 + with Stubber(fs._sync_fs._client) as stubber: + stubber.add_response("head_object", {"ContentLength": size, "VersionId": "v1"}, None) + stubber.add_response( + "copy_object", + {}, + { + "CopySource": {"Bucket": "bucket", "Key": "src", "VersionId": "v1"}, + "Bucket": "bucket", + "Key": "dst", + "ContentType": "text/csv", + "RequestPayer": "requester", + }, + ) + await fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + ContentType="text/csv", + RequestPayer="requester", + ) + stubber.assert_no_pending_responses() @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_copies_source(self): @@ -164,12 +122,6 @@ async def test_copy_object_with_multipart_upload_copies_source(self): # TestS3FileSystem.test_copy_object_with_multipart_upload_copies_source. await self._multipart_copy() - @pytest.mark.asyncio - async def test_copy_object_with_multipart_upload_failed_listing(self): - # GH-973: nothing is written when the annotations cannot be listed. - with pytest.raises(PermissionError): - await self._multipart_copy(fail_list=True) - @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_failed_part(self): # GH-973: a failed part copy aborts the upload, as in S3FileSystem. @@ -278,7 +230,9 @@ async def test_copy_object_with_multipart_upload_max_workers(self): {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} ) ) - call, state = overlapping_calls(3) + # Waits briefly for a fourth overlapping call, which only a copy that + # exceeds max_workers runs. + call, state = overlapping_calls(4, timeout=0.5) def upload_part_copy(**kw): call() @@ -372,12 +326,41 @@ async def test_max_workers_none(self): fs = AioS3FileSystem( connection=mock.MagicMock(), max_workers=None, skip_instance_cache=True ) - fs._sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) + sync_fs = fs._sync_fs + sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) assert await fs._info("s3://bucket/key") == {"name": "bucket/key"} - # The files and copies take the same default. - cpus = getattr(os, "process_cpu_count", os.cpu_count)() - assert fs._sync_fs.max_workers == fs._thread_pool.max_workers == min(32, (cpus or 1) + 4) + # A multipart copy and a file without max_workers of their own. + sync_fs.core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload( + {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} + ) + ) + sync_fs.core.upload_part_copy = mock.MagicMock( + side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + ) + sync_fs.core.complete_multipart_upload = mock.MagicMock() + sync_fs.core.put_object = mock.MagicMock() + # The HeadObject of the source, with 2 parts of the default block size. + sync_fs._call = sync_fs._core.call = mock.MagicMock( + return_value={"ContentLength": 2 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE} + ) + await fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) + assert sync_fs.core.upload_part_copy.call_count == 2 + sync_fs.core.complete_multipart_upload.assert_called_once() + + def write(): + with fs.open("s3://bucket/key", "wb") as f: + f.write(b"a") + + await asyncio.to_thread(write) + sync_fs.core.put_object.assert_called_once() @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio @@ -681,28 +664,6 @@ def record_completion(fn, *args, **kwargs): assert events == ["abort"] - @pytest.mark.asyncio - async def test_cancelled_call_stops_retries(self): - # GH-1203: cancelling a call stops the retries of its request in the - # thread of the filesystem, whose backoff wait would take at least - # EVENT_TIMEOUT seconds. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - operation = ThrottledOperation(block=True, operation_name="HeadObject") - fs._sync_fs.info = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) - task = asyncio.ensure_future(fs._info("s3://bucket/key")) - try: - assert await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - finally: - operation.release.set() - started = time.monotonic() - # Waits for the thread of the cancelled call. - await asyncio.to_thread(fs._thread_pool.shutdown, True) - assert time.monotonic() - started < EVENT_TIMEOUT - assert operation.calls == 1 - @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_failed_part_stops_running_parts(self): # GH-1203: a failed part copy stops the retries of the running ones @@ -844,12 +805,14 @@ async def test_copy_object_with_multipart_upload_failed_annotation_skips_queued_ return_value=SimpleNamespace(version_id=None, etag='"e"') ) + loop = asyncio.get_running_loop() + def copy_object_annotation(name, *args, **kwargs): - # Fails once the other copies wait for the only thread. - deadline = time.monotonic() + EVENT_TIMEOUT - while fs._thread_pool._get_executor()._work_queue.qsize() < 2: - assert time.monotonic() < deadline - time.sleep(0.01) + # The tasks of the three copies start in one pass of the event + # loop, in which each submits its copy; a callback scheduled now + # runs after that pass, so the other copies wait for the only + # thread when this one fails. + asyncio.run_coroutine_threadsafe(asyncio.sleep(0), loop).result(EVENT_TIMEOUT) raise PermissionError("Access Denied") sync_fs.core.copy_object_annotation = mock.MagicMock(side_effect=copy_object_annotation) @@ -1010,28 +973,6 @@ async def main(): assert events == expected - @pytest.mark.parametrize( - "block_size", - [ - S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - 1, - S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + 1, - ], - ) - @pytest.mark.asyncio - async def test_copy_object_with_multipart_upload_invalid_block_size(self, block_size): - # GH-926: the message states the accepted range. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() - - with pytest.raises( - ValueError, - match=r"between 5 MiB \(5242880 bytes\) and 5 GiB \(5368709120 bytes\), inclusive", - ): - await fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), S3Path("bucket", "dst"), block_size=block_size - ) - fs._sync_fs._call.assert_not_called() - @pytest.mark.parametrize("commit", [True, False]) def test_transaction_pipe_put_file(self, tmp_path, commit): # GH-977: pipe_file() and put_file() join the transaction of this @@ -1212,25 +1153,6 @@ def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, kwargs): fs.open.assert_not_called() fs._sync_fs._call.assert_not_called() - def test_transaction_put_file_block_size(self, tmp_path): - # In a transaction, put_file() passes block_size and max_workers to - # open() instead of the S3 API, as outside one. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs.open = mock.MagicMock() - local = tmp_path / "local" - local.write_bytes(b"a" * 13) - - with fs.transaction: - fs.put_file(str(local), "s3://bucket/key", block_size=8) - - fs.open.assert_called_once_with( - "s3://bucket/key", - "wb", - block_size=8, - max_workers=fs._sync_fs.max_workers, - s3_additional_kwargs={}, - ) - def test_touch_sync_wrapper(self): # GH-977: touch() used to be fsspec's open()-based default, which # dropped the PutObject parameters and returned None. @@ -1252,8 +1174,8 @@ def test_touch_sync_wrapper(self): @pytest.mark.asyncio async def test_rm_requests(self): # GH-971: _rm() sent the keys of every bucket to the bucket of the - # first path, dropped the DeleteObjects parameters, ignored per-key - # errors and emptied a bucket path. + # first path and dropped the DeleteObjects parameters. The async + # filesystem sends the request of each bucket in parallel. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) @@ -1267,23 +1189,6 @@ async def test_rm_requests(self): ("b2", [{"Key": "b"}], "111122223333"), ] - sync_fs._call.reset_mock() - with pytest.raises(ValueError, match="Cannot delete the bucket"): - await fs._rm("s3://bucket", recursive=True) - sync_fs._call.assert_not_called() - - sync_fs._call.return_value = { - "Errors": [{"Key": "locked", "Code": "AccessDenied", "Message": "Access Denied"}] - } - with pytest.raises(OSError, match=r"bucket/locked \(AccessDenied: Access Denied\)"): - await fs._rm("s3://bucket/locked") - - sync_fs._call.side_effect = PermissionError("Access Denied") - sync_fs.dircache["bucket/a"] = [] - with pytest.raises(PermissionError, match="Access Denied"): - await fs._rm("s3://bucket/a") - assert "bucket/a" not in sync_fs.dircache - @pytest.mark.asyncio async def test_rm_request_error_keeps_errors(self): fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) @@ -1345,46 +1250,75 @@ async def wait_invalidated(path): @pytest.mark.asyncio async def test_rm_maxdepth(self): - # GH-962: _rm() did not pass maxdepth when expanding the path. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - sync_fs = fs._sync_fs - sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) - sync_fs.find = mock.MagicMock(return_value=["bucket/dir/a"]) - sync_fs.exists = mock.MagicMock(return_value=True) - - # batch_size is part of fsspec's async _rm() signature. - await fs._rm("s3://bucket/dir", recursive=True, maxdepth=1, batch_size=10) - sync_fs.find.assert_called_once_with("bucket/dir", maxdepth=1, withdirs=True, detail=False) - (call,) = sync_fs._call.call_args_list - assert call.kwargs["Delete"]["Objects"] == [{"Key": "dir"}, {"Key": "dir/a"}] + # GH-962: _rm() did not pass maxdepth when expanding the path, so a + # key below maxdepth was deleted too. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + with Stubber(fs._sync_fs._client) as stubber: + # One level of the directory: dir/sub/deep is not listed. + stubber.add_response( + "list_objects_v2", + { + "Contents": [{"Key": "dir/a"}], + "CommonPrefixes": [{"Prefix": "dir/sub/"}], + "KeyCount": 2, + }, + {"Bucket": "bucket", "Prefix": "dir/", "Delimiter": "/"}, + ) + stubber.add_response( + "delete_objects", + {}, + { + "Bucket": "bucket", + "Delete": { + "Objects": [{"Key": "dir"}, {"Key": "dir/a"}, {"Key": "dir/sub"}], + "Quiet": True, + }, + }, + ) + # batch_size is part of fsspec's async _rm() signature. + await fs._rm("s3://bucket/dir", recursive=True, maxdepth=1, batch_size=10) + stubber.assert_no_pending_responses() @pytest.mark.parametrize(("mode", "open_mode"), [("overwrite", "wb"), ("create", "xb")]) def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode): - # GH-969: the open() parameters of put_file() go to open(), and the - # other parameters, also in s3_additional_kwargs, to S3. - # GH-972: fsspec's mode argument selects the mode of the file. + # GH-969: in a transaction, the open() parameters of put_file() go to + # open(), and the other parameters, also in s3_additional_kwargs, to + # S3. GH-972: fsspec's mode argument selects the mode of the file. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) fs.open = mock.MagicMock() lpath = tmp_path / "data.csv" lpath.write_bytes(b"a") - fs._put_file_in_transaction( - str(lpath), - "s3://bucket/key", - Callback(), - mode, - block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, - max_workers=2, - StorageClass="STANDARD_IA", - ) - - fs.open.assert_called_once_with( - "s3://bucket/key", - open_mode, - block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, - max_workers=2, - s3_additional_kwargs={"StorageClass": "STANDARD_IA", "ContentType": "text/csv"}, - ) + with fs.transaction: + fs.put_file( + str(lpath), + "s3://bucket/key", + mode=mode, + block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, + max_workers=2, + StorageClass="STANDARD_IA", + ) + # Without them, the file gets the defaults of the filesystem. + fs.put_file(str(lpath), "s3://bucket/key", mode=mode) + + assert fs.open.call_args_list == [ + mock.call( + "s3://bucket/key", + open_mode, + block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE, + max_workers=2, + s3_additional_kwargs={"StorageClass": "STANDARD_IA", "ContentType": "text/csv"}, + ), + mock.call( + "s3://bucket/key", + open_mode, + block_size=fs._sync_fs.default_block_size, + max_workers=fs._sync_fs.max_workers, + s3_additional_kwargs={"ContentType": "text/csv"}, + ), + ] @pytest.mark.asyncio async def test_cp_file_directory(self): @@ -1410,42 +1344,47 @@ async def test_copy_version(self): "bucket/b?versionId=v1", "s3://bucket/d/b", RequestPayer="requester" ) - @pytest.mark.asyncio - async def test_get_version_path_destination(self, tmp_path): - # A Path destination is paired too: the version is downloaded to the - # file, not into a directory of that name. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs.isdir = mock.MagicMock(return_value=False) - fs._get_file = mock.AsyncMock() - - await fs._get("s3://bucket/b?versionId=v1", tmp_path / "out.bin") - assert fs._get_file.await_args.args[:2] == ( - "bucket/b?versionId=v1", - (tmp_path / "out.bin").as_posix(), - ) - @pytest.mark.asyncio async def test_expand_path_glob_lists_stem_prefix(self): # Unversioned globs keep fsspec's async expansion, which lists only # the keys that start with the stem before the first wildcard. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._find = mock.MagicMock(return_value=[]) + sync_fs = fs._sync_fs + sync_fs._call = sync_fs._core.call = mock.MagicMock( + side_effect=lambda method, **request: {"Contents": [], "KeyCount": 0} + ) with pytest.raises(FileNotFoundError): await fs._expand_path("s3://bucket/reports/2026-*.csv") - assert fs._sync_fs._find.call_args.kwargs["prefix"] == "2026-" + listings = [ + c.kwargs + for c in sync_fs._call.call_args_list + if c.args[0] is sync_fs._client.list_objects_v2 + ] + assert listings[0]["Prefix"] == "reports/2026-" + @pytest.mark.parametrize( + ("destination", "expected"), + [ + # GH-979: a version is downloaded to a local path named after its + # key. + (lambda tmp_path: f"{tmp_path}/", "b"), + # A Path destination is paired too: the version is downloaded to + # the file, not into a directory of that name. + (lambda tmp_path: tmp_path / "out.bin", "out.bin"), + ], + ids=["directory", "path"], + ) @pytest.mark.asyncio - async def test_get_version(self, tmp_path): - # GH-979: a version is downloaded to a local path named after its key. + async def test_get_version(self, tmp_path, destination, expected): fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) fs._sync_fs.isdir = mock.MagicMock(return_value=False) fs._get_file = mock.AsyncMock() - await fs._get("s3://bucket/b?versionId=v1", f"{tmp_path}/") + await fs._get("s3://bucket/b?versionId=v1", destination(tmp_path)) assert fs._get_file.await_args.args[:2] == ( "bucket/b?versionId=v1", - f"{tmp_path.as_posix()}/b", + f"{tmp_path.as_posix()}/{expected}", ) @pytest.mark.asyncio @@ -1506,98 +1445,6 @@ async def copy_file(path1, path2, **kwargs): assert finished == ["s3://bucket/b"] fs._sync_fs._call.assert_not_called() - @pytest.mark.parametrize("status", [None, "Suspended", "Enabled"]) - @pytest.mark.asyncio - async def test_mv_null_version_onto_key(self, status): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock( - return_value={} if status is None else {"Status": status} - ) - fs._copy_file = mock.AsyncMock(return_value=True) - fs._delete_objects = mock.AsyncMock() - sources = ["s3://bucket/a?versionId=null", "s3a://bucket/b?version_id=null"] - destinations = ["s3://bucket/a", "s3://bucket/b"] - - await fs._mv(sources, destinations, MetadataDirective="COPY") - - fs._sync_fs._call.assert_called_once_with( - fs._sync_fs._client.get_bucket_versioning, Bucket="bucket" - ) - if status == "Enabled": - assert fs._copy_file.await_args_list == [ - mock.call(source, dest, MetadataDirective="COPY") - for source, dest in zip(sources, destinations, strict=True) - ] - fs._delete_objects.assert_awaited_once_with(sources) - else: - fs._copy_file.assert_not_awaited() - fs._delete_objects.assert_awaited_once_with([]) - - @pytest.mark.parametrize("status", [None, "Suspended", "Enabled"]) - @pytest.mark.asyncio - async def test_mv_null_version_conflicts_depend_on_bucket_state(self, status): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock( - return_value={} if status is None else {"Status": status} - ) - fs._copy_file = mock.AsyncMock(return_value=True) - fs._delete_objects = mock.AsyncMock() - sources = ["s3://bucket/a", "s3://bucket/b?versionId=null"] - destinations = ["s3a://bucket/b", "s3://bucket/out"] - - with ( - contextlib.nullcontext() - if status == "Enabled" - else pytest.raises(ValueError, match="another path that is moved") - ): - await fs._mv(sources, destinations) - - if status == "Enabled": - assert fs._copy_file.await_args_list == [ - mock.call(source, dest) for source, dest in zip(sources, destinations, strict=True) - ] - fs._delete_objects.assert_awaited_once_with(sources) - else: - fs._copy_file.assert_not_awaited() - fs._delete_objects.assert_not_awaited() - - @pytest.mark.asyncio - async def test_mv_null_version_directory_bucket_does_not_read_bucket_state(self): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() - fs._copy_file = mock.AsyncMock(return_value=True) - fs._delete_objects = mock.AsyncMock() - key = "s3://example--usw2-az1--x-s3/key" - - await fs._mv([f"{key}?versionId=null"], [key]) - - fs._sync_fs._call.assert_not_called() - fs._copy_file.assert_not_awaited() - fs._delete_objects.assert_awaited_once_with([]) - - @pytest.mark.parametrize("stage", ["lookup", "copy"]) - @pytest.mark.asyncio - async def test_mv_null_version_failure_does_not_delete(self, stage): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock( - return_value={"Status": "Enabled"} - ) - fs._copy_file = mock.AsyncMock(return_value=True) - fs._delete_objects = mock.AsyncMock() - if stage == "lookup": - fs._sync_fs._call.side_effect = PermissionError("Access Denied") - else: - fs._copy_file.side_effect = PermissionError("Access Denied") - - with pytest.raises(PermissionError, match="Access Denied"): - await fs._mv( - ["s3://bucket/other", "s3://bucket/key?versionId=null"], - ["s3://bucket/out", "s3://bucket/key"], - ) - fs._delete_objects.assert_not_awaited() - if stage == "lookup": - fs._copy_file.assert_not_awaited() - @pytest.mark.parametrize("size", [10, 5 * 2**30 + 1]) @pytest.mark.asyncio async def test_cp_file_multipart_parameters(self, size): @@ -1693,154 +1540,6 @@ def test_internal_file_system_not_cached(self): for fs in (fs1, fs2, fs3): assert fs._sync_fs not in S3FileSystem._cache.values() - @pytest.mark.parametrize( - ("code", "exception"), - [ - (None, None), - ("NoSuchUpload", None), - ("NoSuchBucket", FileNotFoundError), - ("AccessDenied", PermissionError), - ("InternalError", OSError), - ], - ) - def test_clear_multipart_uploads_race(self, code, exception): - fs = AioS3FileSystem( - key="dummy", - secret="dummy", - region_name="us-east-1", - max_workers=1, - retry_config=RetryConfig(attempt=1), - skip_instance_cache=True, - ) - with Stubber(fs.core.client) as stubber: - stubber.add_response( - "list_multipart_uploads", - { - "Uploads": [ - {"Key": "prefix/gone", "UploadId": "gone"}, - { - "Key": "prefix/pending", - "UploadId": "pending", - "ChecksumAlgorithm": "CRC32", - "ChecksumType": "FULL_OBJECT", - }, - ], - "IsTruncated": False, - }, - {"Bucket": "bucket", "Prefix": "prefix/"}, - ) - request = {"Bucket": "bucket", "Key": "prefix/gone", "UploadId": "gone"} - if code: - stubber.add_client_error( - "abort_multipart_upload", - service_error_code=code, - http_status_code=404 if code.startswith("NoSuch") else 500, - expected_params=request, - ) - else: - stubber.add_response("abort_multipart_upload", {}, request) - stubber.add_response( - "abort_multipart_upload", - {}, - {"Bucket": "bucket", "Key": "prefix/pending", "UploadId": "pending"}, - ) - expected = pytest.raises(exception) if exception else contextlib.nullcontext() - with expected: - fs.clear_multipart_uploads("s3://bucket/prefix/") - stubber.assert_no_pending_responses() - - def test_clear_multipart_uploads_empty(self): - fs = AioS3FileSystem( - key="dummy", - secret="dummy", - region_name="us-east-1", - max_workers=1, - skip_instance_cache=True, - ) - with Stubber(fs.core.client) as stubber: - stubber.add_response( - "list_multipart_uploads", - {"Uploads": [], "IsTruncated": False}, - {"Bucket": "bucket", "Prefix": "prefix/"}, - ) - fs.clear_multipart_uploads("s3://bucket/prefix/") - stubber.assert_no_pending_responses() - - @pytest.mark.parametrize( - ("algorithm", "checksum_type"), - [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], - ) - @pytest.mark.asyncio - async def test_multipart_copy_uses_creation_algorithm(self, algorithm, checksum_type): - fs = AioS3FileSystem( - key="dummy", - secret="dummy", - region_name="us-east-1", - max_workers=1, - skip_instance_cache=True, - ) - block_size = 5 * 2**30 - size = 2 * block_size - checksum_kwargs = {"ChecksumAlgorithm": algorithm} if algorithm else {} - if checksum_type: - checksum_kwargs["ChecksumType"] = checksum_type - expected_parts = [] - with Stubber(fs.core.client) as stubber: - stubber.add_response( - "head_object", {"ContentLength": size}, {"Bucket": "bucket", "Key": "src"} - ) - stubber.add_response( - "create_multipart_upload", - {"Bucket": "bucket", "Key": "dst", "UploadId": "u", **checksum_kwargs}, - {"Bucket": "bucket", "Key": "dst", **checksum_kwargs}, - ) - for number, range_ in ( - (1, f"bytes=0-{block_size - 1}"), - (2, f"bytes={block_size}-{size - 1}"), - ): - result = {"ETag": f'"p{number}"', f"Checksum{algorithm or 'CRC32'}": "checksum"} - stubber.add_response( - "upload_part_copy", - {"CopyPartResult": result}, - { - "Bucket": "bucket", - "Key": "dst", - "UploadId": "u", - "PartNumber": number, - "CopySource": {"Bucket": "bucket", "Key": "src"}, - "CopySourceRange": range_, - }, - ) - expected_parts.append( - { - "ETag": result["ETag"], - "PartNumber": number, - **({f"Checksum{algorithm}": "checksum"} if algorithm else {}), - } - ) - stubber.add_response( - "complete_multipart_upload", - {"ETag": '"done"'}, - { - "Bucket": "bucket", - "Key": "dst", - "UploadId": "u", - "MultipartUpload": {"Parts": expected_parts}, - **({"ChecksumType": checksum_type} if checksum_type else {}), - }, - ) - kwargs = { - "source": S3Path("bucket", "src"), - "destination": S3Path("bucket", "dst"), - "block_size": block_size, - "MetadataDirective": "REPLACE", - "TaggingDirective": "REPLACE", - "AnnotationDirective": "EXCLUDE", - **checksum_kwargs, - } - await fs._copy_object_with_multipart_upload(**kwargs) - stubber.assert_no_pending_responses() - @pytest.fixture(scope="class") def fs(self, request): if not hasattr(request, "param"): @@ -3016,15 +2715,13 @@ def read(): assert sync_fs.core.get_object.call_count == 4 assert state["peak"] == 2 - @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio - async def test_open_parts_in_flight(self, asynchronous): + async def test_open_parts_in_flight(self): # GH-1158: a write keeps at most max_workers parts in flight with the - # executor of either kind. The parts waiting for a worker - # used to be queued with their data. - fs = AioS3FileSystem( - connection=mock.MagicMock(), asynchronous=asynchronous, skip_instance_cache=True - ) + # S3AioExecutor of the file; the S3ThreadPoolExecutor of an + # asynchronous filesystem is that of S3File. The parts waiting for a + # worker used to be queued with their data. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE gates: dict[int, threading.Event] = {} @@ -3071,15 +2768,12 @@ def write(): sync_fs.core.complete_multipart_upload.assert_called_once() assert in_flight == [1, 2, 2, 2] - @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio - async def test_open_failed_part_while_waiting(self, asynchronous): + async def test_open_failed_part_while_waiting(self): # GH-1173: a part found failed when a write waits for a part stops the - # write with its error with the executor of either kind. The write - # used to upload the rest of the data. - fs = AioS3FileSystem( - connection=mock.MagicMock(), asynchronous=asynchronous, skip_instance_cache=True - ) + # write with its error with the S3AioExecutor of the file, as with + # that of S3File. The write used to upload the rest of the data. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( @@ -3126,26 +2820,9 @@ def test_open_invalid_max_workers(self): with pytest.raises(ValueError, match="max_workers must be greater than 0"): fs.open("s3://bucket/key", "wb", max_workers=0) - def test_open_version_id(self): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs.info = mock.MagicMock( - return_value=S3Object( - init={"Key": "key"}, - type=S3ObjectType.S3_OBJECT_TYPE_FILE, - bucket="bucket", - key="key", - ) - ) - fs._sync_fs.info.return_value.size = 4 - - with fs.open("s3://bucket/key", "rb", version_id="v1") as f: - assert isinstance(f, AioS3File) - assert f.version_id == "v1" - assert f.size == 4 - fs._sync_fs.info.assert_called_once_with("bucket/key?versionId=v1", version_id="v1") - def test_open_lookup_parameters(self): - # GH-1004: the lookup of the file sends its lookup parameters. + # GH-1004: the lookup of the file sends its lookup parameters, and + # looks up the version that is opened. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) fs._sync_fs.info = mock.MagicMock( return_value=S3Object( @@ -3157,38 +2834,14 @@ def test_open_lookup_parameters(self): ) sse_c = {"SSECustomerAlgorithm": "AES256", "SSECustomerKey": "k" * 32} - with fs.open("s3://bucket/key", "rb", ContentType="text/plain", **sse_c) as f: + with fs.open( + "s3://bucket/key", "rb", version_id="v1", ContentType="text/plain", **sse_c + ) as f: assert isinstance(f, AioS3File) - fs._sync_fs.info.assert_called_once_with("bucket/key", version_id=None, **sse_c) - - @pytest.mark.parametrize( - ("start", "end", "max_workers", "worker_block_size", "ranges"), - [ - (42, 1337, 1, 999, [(42, 1337)]), # single worker - (42, 1337, 2, 999, [(42, 42 + 999), (42 + 999, 1337)]), # more workers - ( - 42, - 1337, - 2, - 333, - [ - (42, 42 + 333), - (42 + 333, 42 + 666), - (42 + 666, 42 + 999), - (42 + 999, 1337), - ], - ), - (42, 1337, 2, 1295, [(42, 1337)]), # single block - (42, 1337, 2, 1296, [(42, 1337)]), # single block - (42, 1337, 2, 1294, [(42, 1336), (1336, 1337)]), # single block too small - ], - ) - def test_get_ranges(self, start, end, max_workers, worker_block_size, ranges): - assert ( - S3File._get_ranges( - start, end, max_workers=max_workers, worker_block_size=worker_block_size - ) - == ranges + assert f.version_id == "v1" + assert f.size == 4 + fs._sync_fs.info.assert_called_once_with( + "bucket/key?versionId=v1", version_id="v1", **sse_c ) @@ -3248,7 +2901,9 @@ def upload_part(**kw): written = asyncio.Event() async def release(): + deadline = time.monotonic() + EVENT_TIMEOUT while True: + assert time.monotonic() < deadline with lock: if written.is_set() or len(active) >= 2: # Only the newest part finishes until the data is @@ -3270,7 +2925,8 @@ async def release(): await f.close() await releaser - assert in_flight == [1, 2, 2, 2] + assert len(in_flight) == 4 + assert max(in_flight) <= 2 (_, parts), _ = fs.core.complete_multipart_upload.call_args assert [part.part_number for part in parts] == [1, 2, 3, 4] fs.core.abort_multipart_upload.assert_not_called() @@ -3703,43 +3359,29 @@ async def test_cancelled_creation_stops_retries(self, fs): assert operation.calls == 1 assert f.closed - @pytest.mark.asyncio - async def test_discard_in_stopped_scope(self, fs): - # GH-1203: a stop scope does not stop the retries of the abort of - # _discard(). - block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) - await f.write(b"x" * block_size) - await f.flush() - # The abort of the core, whose request is stubbed. - del fs.core.abort_multipart_upload - operation = ThrottledOperation(failures=1, operation_name="AbortMultipartUpload") - fs.core.call = retried(operation, RetryConfig(attempt=3, multiplier=0)) - stop = threading.Event() - stop.set() - - def discard(): - with stop_retries_on(stop): - f._discard() - - await asyncio.to_thread(discard) - assert operation.calls == 2 - # Nothing is left to abort. - await f._close_without_commit() - assert operation.calls == 2 - @pytest.mark.asyncio async def test_cancelled_read_frees_ranges(self, fs): # GH-1186: a cancelled parallel read cancels the reads that have not # finished and frees the data of the others, also while the # cancellation is held, as a failed read does. block_size = 2**10 + loop = asyncio.get_running_loop() read = [] + all_read = threading.Barrier(4, timeout=EVENT_TIMEOUT) + held = threading.Event() + release = threading.Event() def get_object(path, range_, **kwargs): # Unlike bytes, an array can be weakly referenced. data = array.array("B", bytes(range_[1] - range_[0])) read.append(weakref.ref(data)) + all_read.wait() + if range_[0] == 0: + # Held until the read is cancelled, after a pass of the event + # loop that takes the results of the other ranges. + asyncio.run_coroutine_threadsafe(asyncio.sleep(0), loop).result(EVENT_TIMEOUT) + held.set() + assert release.wait(EVENT_TIMEOUT) return data fs.core.get_object = mock.MagicMock(side_effect=get_object) @@ -3751,21 +3393,16 @@ def get_object(path, range_, **kwargs): ) info.size = block_size * 4 fs._sync_fs.info = mock.MagicMock(return_value=info) - wait_reads = asyncio.wait - - async def cancelled_wait(reads, **kwargs): - # As asyncio.wait(), whose frame keeps the tasks in a set, - # cancelled once the reads have finished. - reads = set(reads) - await wait_reads(reads) - raise asyncio.CancelledError f = await fs.open_async("s3://bucket/key", "rb", block_size=block_size, max_workers=4) - with ( - mock.patch.object(asyncio, "wait", side_effect=cancelled_wait), - pytest.raises(asyncio.CancelledError) as info_, - ): - await f.read() + task = asyncio.ensure_future(f.read()) + try: + assert await asyncio.to_thread(held.wait, EVENT_TIMEOUT) + task.cancel() + with pytest.raises(asyncio.CancelledError) as info_: + await task + finally: + release.set() await f.close() await asyncio.to_thread(f._pool.shutdown, True) for _ in range(3): From c9628565228e23cde1fa0351b834d484312196f8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:33:45 +0900 Subject: [PATCH 05/19] Remove the async S3 tests that repeat the sync ones against S3 (GH-1239) Most AioS3FileSystem methods run the S3FileSystem method of the same name in a thread of the pool, and AioS3File differs from S3File only in its executor. Remove the real-S3 tests that repeat the sync tests through those delegations: listing, info, exists, find, du, touch, cat, put, get, rm_file, checksum, sign, mkdir, metadata and tags, the small writes and appends, the checksum multipart writes, the read through core and _fetch_range, and the sync wrapper tests. Check that find() passes maxdepth to the listing in test_find, which test_find_maxdepth checked; glob() still checks withdirs. Keep only the multipart case of test_write_transaction and one size of test_upload_cp_file, and remove 2 files instead of 10 in test_rm_recursive. The fs fixture no longer takes parameters. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 833 +-------------------- 1 file changed, 11 insertions(+), 822 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 5beb2e9f3..b965c9c1c 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -7,14 +7,9 @@ import tempfile import threading import time -import urllib.parse -import urllib.request import uuid import weakref from concurrent.futures import Future, ThreadPoolExecutor, wait -from datetime import UTC, datetime -from itertools import chain -from pathlib import Path from types import SimpleNamespace from unittest import mock @@ -22,7 +17,6 @@ import fsspec import pytest from botocore.stub import Stubber -from fsspec import Callback from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile @@ -33,7 +27,6 @@ S3MultipartUploadPart, S3Object, S3ObjectType, - S3StorageClass, ) from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_writer import S3MultipartStateError @@ -1541,202 +1534,19 @@ def test_internal_file_system_not_cached(self): assert fs._sync_fs not in S3FileSystem._cache.values() @pytest.fixture(scope="class") - def fs(self, request): - if not hasattr(request, "param"): - request.param = {} - return AioS3FileSystem(connection=connect(), **request.param) - - @pytest.mark.parametrize( - ("algorithm", "checksum_type"), - [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], - ) - def test_open_multipart_with_checksum(self, fs, algorithm, checksum_type): - block_size = 5 * 2**20 - data = b"x" * (block_size + 1) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_open_multipart_with_checksum/{uuid.uuid4()}" - ) - kwargs = {"ChecksumAlgorithm": algorithm} if algorithm else {} - if checksum_type: - kwargs["ChecksumType"] = checksum_type - try: - with fs.open(path, "wb", block_size=block_size, s3_additional_kwargs=kwargs) as file: - file.write(data) - assert fs.cat_file(path) == data - assert fs.list_multipart_uploads(path) == [] - finally: - fs.clear_multipart_uploads(path) - if fs.exists(path): - fs.rm(path) - - @pytest.mark.parametrize( - ("algorithm", "checksum_type"), - [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], - ) - def test_put_file_multipart_with_checksum(self, fs, tmp_path, algorithm, checksum_type): - block_size = 5 * 2**20 - data = b"x" * (block_size + 1) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_put_file_multipart_with_checksum/{uuid.uuid4()}" - ) - kwargs = {"ChecksumAlgorithm": algorithm} if algorithm else {} - if checksum_type: - kwargs["ChecksumType"] = checksum_type - try: - lpath = tmp_path / "data" - lpath.write_bytes(data) - fs.put_file(str(lpath), path, block_size=block_size, s3_additional_kwargs=kwargs) - assert fs.cat_file(path) == data - assert fs.list_multipart_uploads(path) == [] - finally: - fs.clear_multipart_uploads(path) - if fs.exists(path): - fs.rm(path) - - @pytest.mark.parametrize( - ("algorithm", "checksum_type"), - [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], - ) - def test_pipe_file_multipart_with_checksum(self, fs, algorithm, checksum_type): - block_size = 5 * 2**20 - data = b"x" * (block_size + 1) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_pipe_file_multipart_with_checksum/{uuid.uuid4()}" - ) - kwargs = {"ChecksumAlgorithm": algorithm} if algorithm else {} - if checksum_type: - kwargs["ChecksumType"] = checksum_type - try: - fs.pipe_file(path, data, block_size=block_size, s3_additional_kwargs=kwargs) - assert fs.cat_file(path) == data - assert fs.list_multipart_uploads(path) == [] - finally: - fs.clear_multipart_uploads(path) - if fs.exists(path): - fs.rm(path) - - @pytest.mark.parametrize( - ("algorithm", "checksum_type"), - [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], - ) - def test_append_multipart_with_checksum(self, fs, algorithm, checksum_type): - block_size = 5 * 2**20 - data = b"x" * (block_size + 1) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_append_multipart_with_checksum/{uuid.uuid4()}" - ) - kwargs = {"ChecksumAlgorithm": algorithm} if algorithm else {} - if checksum_type: - kwargs["ChecksumType"] = checksum_type - try: - original = b"y" * block_size - fs.pipe_file(path, original) - with fs.open(path, "ab", block_size=block_size, s3_additional_kwargs=kwargs) as file: - file.write(data) - assert fs.cat_file(path) == original + data - assert fs.list_multipart_uploads(path) == [] - finally: - fs.clear_multipart_uploads(path) - if fs.exists(path): - fs.rm(path) - - def test_clear_multipart_uploads_after_listed_upload_is_aborted(self, fs): - sync_fs = fs._sync_fs - prefix = ( - f"{ENV.s3_staging_key}{ENV.schema}/filesystem/test_async_clear_multipart_race/" - f"{uuid.uuid4()}/" - ) - path = f"s3://{ENV.s3_staging_bucket}/{prefix}" - gone_path = S3Path(ENV.s3_staging_bucket, f"{prefix}gone") - gone = sync_fs.core.create_multipart_upload(gone_path) - try: - sync_fs.core.create_multipart_upload(S3Path(ENV.s3_staging_bucket, f"{prefix}pending")) - list_uploads = sync_fs.list_multipart_uploads - - def list_then_abort(path): - uploads = list_uploads(path) - assert len(uploads) == 2 - assert any(upload.upload_id == gone.upload_id for upload in uploads) - sync_fs.core.abort_multipart_upload(gone) - return uploads - - with mock.patch.object( - sync_fs, "list_multipart_uploads", side_effect=list_then_abort - ) as list_mock: - fs.clear_multipart_uploads(path) - list_mock.assert_called_once_with(path) - assert fs.list_multipart_uploads(path) == [] - finally: - fs.clear_multipart_uploads(path) - - @pytest.mark.parametrize( - ("fs", "start", "end", "target_data"), - list( - chain( - *[ - [ - ({"default_block_size": x}, 0, 5, b"01234"), - ({"default_block_size": x}, 2, 7, b"23456"), - ({"default_block_size": x}, 0, 10, b"0123456789"), - ] - for x in (S3FileSystem.DEFAULT_BLOCK_SIZE, 3) - ] - ) - ), - indirect=["fs"], - ) - def test_read(self, fs, start, end, target_data): - # lowest level access: use the core - data = fs._sync_fs.core.get_object( - S3Path(ENV.s3_staging_bucket, ENV.s3_filesystem_test_file_key), (start, end) - ) - assert data == target_data, data - with fs.open( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}", "rb" - ) as file: - # mid-level access: use _fetch_range - data = file._fetch_range(start, end) - assert data == target_data, data - # high-level: use fileobj seek and read - file.seek(start) - data = file.read(end - start) - assert data == target_data, data - - @pytest.mark.parametrize( - ("base", "exp"), - [ - (1, 2**10), - (1, 2**20), - ], - ) - def test_write(self, fs, base, exp): - data = b"a" * (base * exp) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_write/{uuid.uuid4()}" - ) - with fs.open(path, "wb") as f: - f.write(data) - with fs.open(path, "rb") as f: - actual = f.read() - assert len(actual) == len(data) - assert actual == data + def fs(self): + return AioS3FileSystem(connection=connect()) @pytest.mark.parametrize( "size", [ - 2**10, # < block size: one-shot PutObject path (the GH-719 regression) 10 * 2**20, # > block size (5 MiB): multipart path via the async executor ], ) def test_write_transaction(self, fs, size): - # GH-719 regression for the async filesystem: AioS3File inherits - # _upload_chunk/commit from S3File, so the transaction fix must hold here - # too, including the multipart path driven by the async executor. + # GH-719 regression for the async filesystem: the deferred commit of a + # multipart write whose parts the async executor uploads. The + # single-request path is S3File's, as in the sync test. data = b"a" * size path = ( f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" @@ -1749,50 +1559,6 @@ def test_write_transaction(self, fs, size): assert len(actual) == len(data) assert actual == data - def test_write_transaction_rollback(self, fs): - # Kept small-only on purpose: the multipart discard()/abort path is - # already exercised by the sync test (AioS3File inherits discard()), - # so there is no need to pay for another large upload here. - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_write_transaction_rollback/{uuid.uuid4()}" - ) - - def write_then_fail(): - with fs.transaction: - f = fs.open(path, "wb") - f.write(b"hello world") - f.close() - raise RuntimeError("rollback") - - with pytest.raises(RuntimeError): - write_then_fail() - fs.invalidate_cache(path) - assert not fs.exists(path) - - @pytest.mark.parametrize( - ("base", "exp"), - [ - (1, 2**10), - (1, 2**20), - ], - ) - def test_append(self, fs, base, exp): - data = b"a" * (base * exp) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_append/{uuid.uuid4()}" - ) - with fs.open(path, "ab") as f: - f.write(data) - extra = b"extra" - with fs.open(path, "ab") as f: - f.write(extra) - with fs.open(path, "rb") as f: - actual = f.read() - assert len(actual) == len(data + extra) - assert actual == data + extra - @pytest.mark.parametrize( "size", [ @@ -1830,141 +1596,6 @@ async def test_open_async_append_exclusive(self, fs): async with await fs.open_async(path, "rb") as f: assert await f.read() == b"hello world" - @pytest.mark.asyncio - async def test_ls_buckets(self, fs): - fs.invalidate_cache() - actual = await fs._ls("s3://") - assert ENV.s3_staging_bucket in actual, actual - - fs.invalidate_cache() - actual = await fs._ls("s3:///") - assert ENV.s3_staging_bucket in actual, actual - - fs.invalidate_cache() - actual = await fs._ls("s3://", detail=True) - found = next(filter(lambda x: x.name == ENV.s3_staging_bucket, actual), None) - assert found - assert found.name == ENV.s3_staging_bucket - - fs.invalidate_cache() - actual = await fs._ls("s3:///", detail=True) - found = next(filter(lambda x: x.name == ENV.s3_staging_bucket, actual), None) - assert found - assert found.name == ENV.s3_staging_bucket - - @pytest.mark.asyncio - async def test_ls_dirs(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_ls_dirs" - ) - for i in range(5): - await fs._pipe_file(f"{dir_}/prefix/test_{i}", bytes(i)) - await fs._touch(f"{dir_}/prefix2") - - assert len(await fs._ls(f"{dir_}/prefix")) == 5 - assert len(await fs._ls(f"{dir_}/prefix/")) == 5 - assert len(await fs._ls(f"{dir_}/prefix/test_")) == 0 - assert len(await fs._ls(f"{dir_}/prefix2")) == 1 - - test_1 = await fs._ls(f"{dir_}/prefix/test_1") - assert len(test_1) == 1 - assert test_1[0] == fs._strip_protocol(f"{dir_}/prefix/test_1") - - test_1_detail = await fs._ls(f"{dir_}/prefix/test_1", detail=True) - assert len(test_1_detail) == 1 - assert test_1_detail[0].name == fs._strip_protocol(f"{dir_}/prefix/test_1") - assert test_1_detail[0].size == 1 - - @pytest.mark.asyncio - async def test_info_bucket(self, fs): - dir_ = f"s3://{ENV.s3_staging_bucket}" - s3_path = S3Path.parse(dir_) - info = await fs._info(dir_) - - assert info.name == fs._strip_protocol(dir_) - assert info.bucket == s3_path.bucket - assert info.key is None - assert info.last_modified is None - assert info.size == 0 - assert info.etag is None - assert info.type == S3ObjectType.S3_OBJECT_TYPE_DIRECTORY - assert info.storage_class == S3StorageClass.S3_STORAGE_CLASS_BUCKET - assert info.version_id == s3_path.version_id - - dir_ = f"s3://{ENV.s3_staging_bucket}/" - s3_path = S3Path.parse(dir_) - info = await fs._info(dir_) - - assert info.name == fs._strip_protocol(dir_) - assert info.bucket == s3_path.bucket - assert info.key is None - assert info.last_modified is None - assert info.size == 0 - assert info.etag is None - assert info.type == S3ObjectType.S3_OBJECT_TYPE_DIRECTORY - assert info.storage_class == S3StorageClass.S3_STORAGE_CLASS_BUCKET - assert info.version_id == s3_path.version_id - - @pytest.mark.asyncio - async def test_info_dir(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_info_dir" - ) - file = f"{dir_}/{uuid.uuid4()}" - - fs.invalidate_cache() - with pytest.raises(FileNotFoundError): - await fs._info(f"s3://{uuid.uuid4()}") - - await fs._pipe_file(file, b"a") - s3_path = S3Path.parse(dir_) - fs.invalidate_cache() - info = await fs._info(dir_) - fs.invalidate_cache() - - assert info.name == fs._strip_protocol(dir_) - assert info.bucket == s3_path.bucket - assert info.key == s3_path.key.rstrip("/") - assert info.last_modified is None - assert info.size == 0 - assert info.etag is None - assert info.type == S3ObjectType.S3_OBJECT_TYPE_DIRECTORY - assert info.storage_class == S3StorageClass.S3_STORAGE_CLASS_DIRECTORY - assert info.version_id == s3_path.version_id - - @pytest.mark.asyncio - async def test_info_file(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_info_file" - ) - file = f"{dir_}/{uuid.uuid4()}" - - fs.invalidate_cache() - with pytest.raises(FileNotFoundError): - await fs._info(file) - - now = datetime.now(UTC) - await fs._pipe_file(file, b"a") - s3_path = S3Path.parse(file) - fs.invalidate_cache() - info = await fs._info(file) - fs.invalidate_cache() - ls_info = (await fs._ls(file, detail=True))[0] - - assert info == ls_info - assert info.name == fs._strip_protocol(file) - assert info.bucket == s3_path.bucket - assert info.key == s3_path.key - assert info.last_modified >= now - assert info.size == 1 - assert info.etag is not None - assert info.type == S3ObjectType.S3_OBJECT_TYPE_FILE - assert info.storage_class == S3StorageClass.S3_STORAGE_CLASS_STANDARD - assert info.version_id == s3_path.version_id - @pytest.mark.asyncio async def test_find(self, fs): dir_ = ( @@ -1987,6 +1618,10 @@ async def test_find(self, fs): result = await fs._find(f"{dir_}/prefix/test_") assert len(result) == 0 + # maxdepth reaches the listing: only prefix2 is one level down. + result = await fs._find(dir_, maxdepth=1) + assert result == [fs._strip_protocol(f"{dir_}/prefix2")] + result = await fs._find(f"{dir_}/prefix", prefix="test_") assert len(result) == 5 @@ -2004,95 +1639,6 @@ async def test_find(self, fs): ].name == fs._strip_protocol(f"{dir_}/prefix/test_1") assert test_1_detail[fs._strip_protocol(f"{dir_}/prefix/test_1")].size == 1 - @pytest.mark.asyncio - async def test_find_maxdepth(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_find_maxdepth" - ) - # Create files at different depths - await fs._touch(f"{dir_}/file0.txt") - await fs._touch(f"{dir_}/level1/file1.txt") - await fs._touch(f"{dir_}/level1/level2/file2.txt") - await fs._touch(f"{dir_}/level1/level2/level3/file3.txt") - - # maxdepth must be at least 1, as in fsspec - with pytest.raises(ValueError, match="maxdepth must be at least 1"): - await fs._find(dir_, maxdepth=0) - - # Test maxdepth=1 (only files in the root) - result = await fs._find(dir_, maxdepth=1) - assert len(result) == 1 - assert fs._strip_protocol(f"{dir_}/file0.txt") in result - - # Test maxdepth=2 (files in root and level1) - result = await fs._find(dir_, maxdepth=2) - assert len(result) == 2 - assert fs._strip_protocol(f"{dir_}/file0.txt") in result - assert fs._strip_protocol(f"{dir_}/level1/file1.txt") in result - - # Test maxdepth=3 (files in root, level1, and level2) - result = await fs._find(dir_, maxdepth=3) - assert len(result) == 3 - assert fs._strip_protocol(f"{dir_}/level1/level2/file2.txt") in result - - # Test no maxdepth (all files) - result = await fs._find(dir_) - assert len(result) == 4 - - @pytest.mark.asyncio - async def test_find_withdirs(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_find_withdirs" - ) - # Create directory structure with files - await fs._touch(f"{dir_}/file1.txt") - await fs._touch(f"{dir_}/subdir1/file2.txt") - await fs._touch(f"{dir_}/subdir1/subdir2/file3.txt") - await fs._touch(f"{dir_}/subdir3/file4.txt") - - # Test default behavior (withdirs=False) - result = await fs._find(dir_) - assert len(result) == 4 # Only files - for r in result: - assert r.endswith(".txt") - - # Test withdirs=True - result = await fs._find(dir_, withdirs=True) - assert len(result) > 4 # Files and directories - - # Verify directories are included - dirs = [r for r in result if not r.endswith(".txt")] - assert len(dirs) > 0 - assert any("subdir1" in d for d in dirs) - assert any("subdir2" in d for d in dirs) - assert any("subdir3" in d for d in dirs) - - # Test withdirs=False explicitly - result = await fs._find(dir_, withdirs=False) - assert len(result) == 4 # Only files - - @pytest.mark.asyncio - async def test_du(self, fs): - """Disk usage reports file sizes, their total, and the requested depth.""" - directory = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/filesystem/test_async_du" - first = f"{directory}/first" - second = f"{directory}/nested/second" - try: - await fs._pipe_file(first, b"abc") - await fs._pipe_file(second, b"12345") - assert await fs._du(directory) == 8 - assert await fs._du(directory, total=False) == { - fs._strip_protocol(first): 3, - fs._strip_protocol(second): 5, - } - assert await fs._du(directory, maxdepth=1) == 3 - assert await fs._du(first) == 3 - finally: - with contextlib.suppress(FileNotFoundError): - await fs._rm(directory, recursive=True) - @pytest.mark.asyncio async def test_glob(self, fs): dir_ = ( @@ -2114,54 +1660,6 @@ async def test_glob(self, fs): with pytest.raises(ValueError): # noqa: PT011 fs.glob("*") - @pytest.mark.asyncio - async def test_exists_bucket(self, fs): - assert await fs._exists("s3://") - assert await fs._exists("s3:///") - - path = f"s3://{ENV.s3_staging_bucket}" - assert await fs._exists(path) - - not_exists_path = f"s3://{uuid.uuid4()}" - assert not await fs._exists(not_exists_path) - - @pytest.mark.asyncio - async def test_exists_object(self, fs): - path = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}" - assert await fs._exists(path) - - not_exists_path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_exists/{uuid.uuid4()}" - ) - assert not await fs._exists(not_exists_path) - - @pytest.mark.asyncio - async def test_rm_file(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_rm_file" - ) - file = f"{dir_}/{uuid.uuid4()}" - await fs._touch(file) - await fs._rm_file(file) - - assert not await fs._exists(file) - assert not await fs._exists(dir_) - - @pytest.mark.asyncio - async def test_rm(self, fs): - dir_ = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_rm" - ) - file = f"{dir_}/{uuid.uuid4()}" - await fs._touch(file) - await fs._rm(file) - - assert not await fs._exists(file) - assert not await fs._exists(dir_) - @pytest.mark.asyncio async def test_rm_recursive(self, fs): dir_ = ( @@ -2169,7 +1667,7 @@ async def test_rm_recursive(self, fs): f"filesystem/test_async_rm_recursive" ) - files = [f"{dir_}/{uuid.uuid4()}" for _ in range(10)] + files = [f"{dir_}/{uuid.uuid4()}" for _ in range(2)] for f in files: await fs._touch(f) @@ -2183,120 +1681,12 @@ async def test_rm_recursive(self, fs): assert not await fs._exists(f) assert not await fs._exists(dir_) - @pytest.mark.asyncio - async def test_touch(self, fs): - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_touch/{uuid.uuid4()}" - ) - assert not await fs._exists(path) - await fs._touch(path) - assert await fs._exists(path) - info = await fs._info(path) - assert info.size == 0 - - with fs.open(path, "wb") as f: - f.write(b"data") - info = await fs._info(path, refresh=True) - assert info.size == 4 - await fs._touch(path, truncate=True) - info = await fs._info(path, refresh=True) - assert info.size == 0 - - with fs.open(path, "wb") as f: - f.write(b"data") - info = await fs._info(path, refresh=True) - assert info.size == 4 - with pytest.raises(ValueError, match="Cannot touch"): - await fs._touch(path, truncate=False) - info = await fs._info(path, refresh=True) - assert info.size == 4 - - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("base", "exp"), - [ - (1, 2**10), - (1, 2**20), - ], - ) - async def test_pipe_cat(self, fs, base, exp): - data = b"a" * (base * exp) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_pipe_cat/{uuid.uuid4()}" - ) - await fs._pipe_file(path, data) - assert await fs._cat_file(path) == data - - @pytest.mark.asyncio - async def test_cat_ranges(self, fs): - data = b"1234567890abcdefghijklmnopqrstuvwxyz" - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_cat_ranges/{uuid.uuid4()}" - ) - await fs._pipe_file(path, data) - - assert await fs._cat_file(path) == data - assert await fs._cat_file(path, start=5) == data[5:] - assert await fs._cat_file(path, end=5) == data[:5] - assert await fs._cat_file(path, start=1, end=-1) == data[1:-1] - assert await fs._cat_file(path, start=-5) == data[-5:] - - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("base", "exp"), - [ - (1, 2**10), - (1, 2**20), - ], - ) - async def test_put(self, fs, base, exp): - with tempfile.NamedTemporaryFile(delete=False) as tmp: - data = b"a" * (base * exp) - tmp.write(data) - tmp.flush() - - rpath = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_put/{uuid.uuid4()}" - ) - await fs._put_file(lpath=tmp.name, rpath=rpath) - tmp.seek(0) - assert await fs._cat_file(rpath) == tmp.read() - @pytest.mark.asyncio @pytest.mark.parametrize( ("base", "exp"), [ + # Below the maximum part size, a single CopyObject request. (1, 2**10), - (1, 2**20), - ], - ) - async def test_put_with_callback(self, fs, base, exp): - with tempfile.NamedTemporaryFile(delete=False) as tmp: - data = b"a" * (base * exp) - tmp.write(data) - tmp.flush() - - rpath = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_put_with_callback/{uuid.uuid4()}" - ) - callback = Callback() - await fs._put_file(lpath=tmp.name, rpath=rpath, callback=callback) - tmp.seek(0) - assert await fs._cat_file(rpath) == tmp.read() - assert callback.size == os.stat(tmp.name).st_size - assert callback.value == callback.size - - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("base", "exp"), - [ - (1, 2**10), - (1, 2**20), ], ) async def test_upload_cp_file(self, fs, base, exp): @@ -2364,65 +1754,6 @@ async def test_move_null_version_onto_key(self, fs, versioning_buckets, status): else: assert after == before - @pytest.mark.skipif( - os.getenv("AWS_ATHENA_S3_VERSIONING_TESTS") != "1", - reason="Set AWS_ATHENA_S3_VERSIONING_TESTS=1 to create versioning test buckets.", - ) - @pytest.mark.parametrize( - "status", - [ - pytest.param(None, id="async-wrapper-None"), - pytest.param("Enabled", id="async-wrapper-Enabled"), - pytest.param("Suspended", id="async-wrapper-Suspended"), - ], - ) - def test_move_null_version_onto_key_sync_wrapper(self, fs, versioning_buckets, status): - client, buckets = versioning_buckets - bucket = buckets[status] - key = "async-wrapper" - path = f"s3://{bucket}/{key}" - before = [ - v - for v in client.list_object_versions(Bucket=bucket, Prefix=key)["Versions"] - if v["Key"] == key - ] - assert any(v["VersionId"] == "null" for v in before) - if status: - assert not next(v for v in before if v["VersionId"] == "null")["IsLatest"] - - fs.mv(f"{path}?versionId=null", path) - - with client.get_object(Bucket=bucket, Key=key)["Body"] as body: - assert body.read() == (b"original" if status != "Suspended" else b"current") - after = [ - v - for v in client.list_object_versions(Bucket=bucket, Prefix=key)["Versions"] - if v["Key"] == key - ] - if status == "Enabled": - assert not any(v["VersionId"] == "null" for v in after) - assert len(after) == len(before) - latest = next(v for v in after if v["IsLatest"]) - assert latest["VersionId"] not in {v["VersionId"] for v in before} - else: - assert after == before - - @pytest.mark.asyncio - async def test_move(self, fs): - path1 = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_move/{uuid.uuid4()}" - ) - path2 = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_move/{uuid.uuid4()}" - ) - data = b"a" - await fs._pipe_file(path1, data) - fs.mv(path1, path2) - assert await fs._cat_file(path2) == data - assert not await fs._exists(path1) - @pytest.mark.asyncio async def test_move_recursive(self, fs): # GH-974: the directory entries used to make mv() fail, after copying @@ -2439,98 +1770,6 @@ async def test_move_recursive(self, fs): assert not await fs._exists(f"{base}/src/a") assert not await fs._exists(f"{base}/src/sub/b") - @pytest.mark.asyncio - async def test_get_file(self, fs): - with tempfile.TemporaryDirectory() as tmp: - rpath = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}" - lpath = Path(f"{tmp}/{uuid.uuid4()}") - callback = Callback() - await fs._get_file(rpath=rpath, lpath=str(lpath), callback=callback) - - assert lpath.open("rb").read() == await fs._cat_file(rpath) - assert callback.size == os.stat(lpath).st_size - assert callback.value == callback.size - - def test_open_returns_aio_s3_file(self, fs): - path = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}" - with fs.open(path, "rb") as f: - assert isinstance(f, AioS3File) - data = f.read() - assert data == b"0123456789" - - def test_checksum(self, fs): - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_checksum/{uuid.uuid4()}" - ) - s3_path = S3Path.parse(path) - - fs.pipe_file(path, b"foo") - checksum = fs.checksum(path) - fs.ls(path) # caching - fs._sync_fs.core.put_object(s3_path, b"bar") - assert checksum == fs.checksum(path) - assert checksum != fs.checksum(path, refresh=True) - - fs.pipe_file(path, b"foo") - checksum = fs.checksum(path) - fs.ls(path) # caching - fs.core.delete_object(s3_path) - assert checksum == fs.checksum(path) - with pytest.raises(FileNotFoundError): - fs.checksum(path, refresh=True) - - def test_sign(self, fs): - path = f"s3://{ENV.s3_staging_bucket}/{ENV.s3_filesystem_test_file_key}" - requested = time.time() - time.sleep(1) - url = fs.sign(path, expiration=100) - parsed = urllib.parse.urlparse(url) - query = urllib.parse.parse_qs(parsed.query) - expires = int(query["Expires"][0]) - with urllib.request.urlopen(url) as r: - data = r.read() - - assert "https" in url - assert requested + 100 < expires - assert data == b"0123456789" - - @pytest.mark.asyncio - async def test_mkdir_and_rmdir(self, fs): - # The bucket already exists. - with pytest.raises(FileExistsError): - await fs._mkdir(f"s3://{ENV.s3_staging_bucket}") - await fs._makedirs(f"s3://{ENV.s3_staging_bucket}", exist_ok=True) - - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_rmdir/{uuid.uuid4()}" - ) - await fs._pipe_file(path, b"data") - # Only bucket paths can be removed. - with pytest.raises(FileExistsError): - await fs._rmdir(path) - with pytest.raises(FileNotFoundError): - fs.rmdir(f"{path}/nonexistent") - - @pytest.mark.asyncio - async def test_metadata_and_tags(self, fs): - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_metadata/{uuid.uuid4()}" - ) - await fs._pipe_file(path, b"data") - - assert fs.metadata(path) == {} - fs.setxattr(path, attr1="value1") - assert fs.metadata(path) == {"attr1": "value1"} - assert fs.getxattr(path, "attr1") == "value1" - assert fs.getxattr(path, "missing") is None - - assert fs.get_tags(path) == {} - fs.put_tags(path, {"tag1": "value1"}) - assert fs.get_tags(path) == {"tag1": "value1"} - def test_pandas_read_csv(self): import pandas @@ -2541,56 +1780,6 @@ def test_pandas_read_csv(self): ) assert [(row["col"],) for _, row in df.iterrows()] == [(123456789,)] - @pytest.mark.parametrize( - "line_count", - [ - 1 * 2**20, - ], - ) - def test_pandas_write_csv(self, line_count): - import pandas - - with tempfile.NamedTemporaryFile("w+t") as tmp: - tmp.write("col1") - tmp.write("\n") - for _ in range(line_count): - tmp.write("a") - tmp.write("\n") - tmp.flush() - - tmp.seek(0) - df = pandas.read_csv(tmp.name) - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_pandas_write_csv/{uuid.uuid4()}.csv" - ) - df.to_csv(path, index=False) - - actual = pandas.read_csv(path) - pandas.testing.assert_frame_equal(actual, df) - - def test_sync_wrappers(self, fs): - """Verify that mirror_sync_methods generates working sync wrappers.""" - actual = fs.ls(f"s3://{ENV.s3_staging_bucket}") - assert isinstance(actual, list) - assert len(actual) > 0 - - assert fs.exists(f"s3://{ENV.s3_staging_bucket}") - - def test_invalidate_cache(self, fs): - path = ( - f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" - f"filesystem/test_async_invalidate_cache/{uuid.uuid4()}" - ) - fs.pipe_file(path, b"data") - fs.info(path) - - # Cache should be populated - fs.invalidate_cache(path) - # Should not raise after cache invalidation - info = fs.info(path) - assert info.size == 4 - class TestAioS3File: @pytest.mark.parametrize("max_workers", [2, None]) From 907fecba15d8e4a0c017a3de5875859f5c421ea5 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:57:05 +0900 Subject: [PATCH 06/19] Check that a call parameter takes precedence over request_kwargs (GH-1239) S3Core.call() documents that a parameter of the request takes precedence over the core's request_kwargs, and no test failed when the order was swapped. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_core.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index ef33479c1..23e9046c8 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -230,15 +230,29 @@ def test_head_bucket(self): ], ) def test_get_object(self, version_id, range_, expected): - core, stubber = _make_core(request_kwargs={"ServerSideEncryption": "AES256"}) - # ServerSideEncryption, which GetObject does not accept, is not sent. + core, stubber = _make_core( + request_kwargs={"ServerSideEncryption": "AES256", "ExpectedBucketOwner": "111122223333"} + ) + # ServerSideEncryption, which GetObject does not accept, is not sent, + # and a parameter of the call takes precedence over request_kwargs. stubber.add_response( "get_object", {"Body": StreamingBody(io.BytesIO(b"data"), 4)}, - {"Bucket": "bucket", "Key": "key", "IfMatch": '"e"', **expected}, + { + "Bucket": "bucket", + "Key": "key", + "IfMatch": '"e"', + "ExpectedBucketOwner": "444455556666", + **expected, + }, ) with stubber: - data = core.get_object(S3Path("bucket", "key", version_id), range_, IfMatch='"e"') + data = core.get_object( + S3Path("bucket", "key", version_id), + range_, + IfMatch='"e"', + ExpectedBucketOwner="444455556666", + ) stubber.assert_no_pending_responses() assert data == b"data" From 6acd2724a639c2d0746c591f26a1347f21277f5d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 23:03:22 +0900 Subject: [PATCH 07/19] Remove and rewrite redundant S3 executor tests (GH-1239) Apply the test_s3_executor.py rows of #1241: remove the tests that restate constructor defaults and ThreadPoolExecutor semantics, the private-callback poll of wait_shielded(), and the dispatch-failure copies; merge the loop shutdown of submit() into test_submit_to_loop_shutdown; keep test_submit_to_executor_shut_down, the only test of the documented error that submit_to() sets when the executor drops a callable, with a bounded wait. Rewrite the cancellation tests without sleeps or private queues: the cancellation after the call has returned now comes from a loop callback that blocks until the call returns, and the cancellation before the start waits for a spy on submit(). test_fork runs a real forked child that submits to the pool. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_executor.py | 276 +++++------------- 1 file changed, 81 insertions(+), 195 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 601d0d048..264aa856e 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -9,11 +9,13 @@ import contextvars import gc import os +import signal import threading import time import traceback import weakref from concurrent.futures import Future, ThreadPoolExecutor, wait +from typing import cast from unittest import mock import pytest @@ -25,13 +27,6 @@ class TestS3AioThreadPool: - def test_init(self): - # None takes the default of ThreadPoolExecutor. - cpus = getattr(os, "process_cpu_count", os.cpu_count)() - assert S3AioThreadPool(None).max_workers == min(32, (cpus or 1) + 4) - with pytest.raises(ValueError, match="max_workers must be greater than 0"): - S3AioThreadPool(0) - def test_run_beyond_default_executor(self): # GH-1159: the calls run in threads of the pool, so that max_workers, # not the default executor of the event loop, bounds them. @@ -101,11 +96,30 @@ def test_run_cancelled_after_return(self): # not stop the work that the call left running, such as the parts of a # write in a transaction. operation = ThrottledOperation(block=True, failures=1) - returned = threading.Event() left = [] + go = threading.Event() pool = S3AioThreadPool(1) + submit = pool.submit + + def submit_and_cancel_after_return(fn, *args, **kwargs): + future = submit(fn, *args, **kwargs) + task = cast(asyncio.Task[None], asyncio.current_task()) + + def cancel_after_return(): + # Runs once the task waits for the call, and blocks the event + # loop until the call has returned, so that the task is + # cancelled before it resumes. + go.set() + future.result(EVENT_TIMEOUT) + task.cancel() + + asyncio.get_running_loop().call_soon(cancel_after_return) + return future + + pool.submit = submit_and_cancel_after_return # type: ignore[method-assign] def work(): + assert go.wait(EVENT_TIMEOUT) context = contextvars.copy_context() thread = threading.Thread( target=context.run, @@ -113,17 +127,9 @@ def work(): ) thread.start() left.append(thread) - returned.set() async def main(): task = asyncio.ensure_future(pool.run(work)) - await asyncio.sleep(0) - # Blocks the event loop, so that the task is cancelled before it - # resumes after the call has returned, which follows the event at - # once. - assert returned.wait(EVENT_TIMEOUT) - time.sleep(0.1) - task.cancel() with pytest.raises(asyncio.CancelledError): await task @@ -150,14 +156,23 @@ def work(): started.set() release.wait(EVENT_TIMEOUT) + running = pool.submit(work) + submit = pool.submit + async def main(): - running = pool.submit(work) + submitted = asyncio.Event() + + def submit_and_notify(fn, *args, **kwargs): + future = submit(fn, *args, **kwargs) + submitted.set() + return future + + pool.submit = submit_and_notify # type: ignore[method-assign] run = pool.run_shielded if shielded else pool.run pending = asyncio.ensure_future(run(events.append, "pending finished")) - await asyncio.to_thread(started.wait, EVENT_TIMEOUT) - # Lets the task submit the call, which waits for the thread. - while not pending.done() and not pool._get_executor()._work_queue.qsize(): - await asyncio.sleep(0.01) + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + # The call waits for the only thread. + await asyncio.wait_for(submitted.wait(), EVENT_TIMEOUT) pending.cancel() with pytest.raises(asyncio.CancelledError): await pending @@ -169,109 +184,39 @@ async def main(): assert events == (["pending finished"] if shielded else []) - def test_wait_shielded_cancelled(self): - # GH-1186: cancelling the waiting task neither cancels the call nor - # lets asyncio log its failure, which ends after the cancellation; - # the caller can still wait for its future. - started = threading.Event() - release = threading.Event() - pool = S3AioThreadPool(1) - - def work(): - started.set() - # The finally block of the test always releases it. - release.wait() - raise PermissionError("Access Denied") - - async def main(): - logged = [] - asyncio.get_running_loop().set_exception_handler( - lambda _, context: logged.append(context["message"]) - ) - running = pool.submit(work) - try: - waiting = asyncio.ensure_future(pool.wait_shielded(running)) - await asyncio.to_thread(started.wait, EVENT_TIMEOUT) - # Lets the task start waiting before it is cancelled. - deadline = time.monotonic() + EVENT_TIMEOUT - while not running._done_callbacks: - assert time.monotonic() < deadline - await asyncio.sleep(0) - waiting.cancel() - with pytest.raises(asyncio.CancelledError): - await waiting - finally: - release.set() - # The worker has run the callbacks of the call once it is - # shut down, and the event loop runs theirs on the next - # iterations. - await asyncio.to_thread(pool.shutdown, True) - with pytest.raises(PermissionError, match="Access Denied"): - running.result() - # The callbacks of the future of the call keep its waiter. - del waiting, running - gc.collect() - for _ in range(3): - await asyncio.sleep(0) - return logged - - assert asyncio.run(main()) == [] - - def test_shutdown(self): - # The calls submitted before the shutdown still run. - events = [] - release = threading.Event() - pool = S3AioThreadPool(1) - pool.submit(release.wait, EVENT_TIMEOUT) - pending = pool.submit(events.append, "pending finished") - pool.shutdown(wait=False) - with pytest.raises(RuntimeError, match="after shutdown"): - pool.submit(sum, [1, 2]) - release.set() - pending.result(timeout=EVENT_TIMEOUT) - - assert events == ["pending finished"] - + @pytest.mark.skipif(not hasattr(os, "fork"), reason="os.fork() is not available.") + @pytest.mark.filterwarnings("ignore:.*fork\\(\\) may lead to deadlocks:DeprecationWarning") def test_fork(self): # A forked child, whose copy of the pool has no threads, creates new - # ones, also if a thread of the parent held the lock of the pool. + # ones. pool = S3AioThreadPool(1) - parent = pool._get_executor() - pool._lock.acquire() - # As os.register_at_fork() runs it in a forked child. - S3AioThreadPool._reset_after_fork() - # The copied executor is kept, so that its collection does not touch - # the locks that the parent's threads may have held. - assert parent in S3AioThreadPool._forked_executors - child = pool._get_executor() assert pool.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 - pool.shutdown() - assert child is not parent - S3AioThreadPool._forked_executors.remove(parent) - parent.shutdown() + pid = os.fork() + if pid == 0: + # The child exits here, whatever happens, instead of running pytest. + code = 1 + try: + if pool.submit(sum, [3, 4]).result(timeout=EVENT_TIMEOUT) == 7: + code = 0 + finally: + os._exit(code) + status = (0, 0) + try: + # A child whose pool waits for the threads of the parent would + # not exit. + deadline = time.monotonic() + 2 * EVENT_TIMEOUT + while (status := os.waitpid(pid, os.WNOHANG)) == (0, 0): + assert time.monotonic() < deadline, "the child did not exit" + time.sleep(0.05) + finally: + if status == (0, 0): + os.kill(pid, signal.SIGKILL) + os.waitpid(pid, 0) + pool.shutdown() + assert os.waitstatus_to_exitcode(status[1]) == 0 class TestS3AioExecutor: - def test_init(self): - # max_workers is optional, as before it was added. - S3AioExecutor(loop=None) - with pytest.raises(ValueError, match="max_workers must be greater than 0"): - S3AioExecutor(loop=None, max_workers=0) - - def test_submit(self): - async def main(): - executor = S3AioExecutor(loop=asyncio.get_running_loop()) - succeeded = executor.submit(sum, [1, 2]) - failed = executor.submit(int, "x") - await asyncio.to_thread(wait, [succeeded, failed]) - return succeeded, failed - - succeeded, failed = asyncio.run(main()) - - assert succeeded.result() == 3 - with pytest.raises(ValueError, match="invalid literal"): - failed.result() - def test_submit_context(self): # GH-1203: a function runs with the context of its submitter, as the # task that dispatches it does, so that the stop scope of the @@ -492,54 +437,6 @@ async def main(): assert not running.cancelled() assert pending.cancelled() - def test_loop_shutdown(self): - # A function that has not started when the event loop shuts down is - # never run, and its future is cancelled and settled, so that wait() - # returns, instead of left pending. - events = [] - started = threading.Event() - settled = threading.Event() - - def work(): - started.set() - # Running until the pending future is settled, so that the - # pending function cannot start before the shutdown. - settled.wait(5) - events.append("finished") - - async def main(): - # The pending function waits for the only thread of the executor. - executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) - running = executor.submit(work) - pending = executor.submit(events.append, "pending finished") - pending.add_done_callback(lambda _: settled.set()) - while not started.is_set(): - await asyncio.sleep(0.01) - return running, pending - - running, pending = asyncio.run(main()) - - _, not_done = wait([running, pending], timeout=5) - assert not not_done - assert events == ["finished"] - assert not running.cancelled() - assert pending.cancelled() - - def test_executor_shut_down(self): - # An error raised before the function starts is set on the future. - async def main(): - executor = S3AioExecutor(loop=asyncio.get_running_loop()) - # Fails the dispatch of the function to the pool. - executor._pool.shutdown() - future = executor.submit(sum, [1, 2]) - with pytest.raises(RuntimeError, match="after shutdown"): - await asyncio.wrap_future(future) - return future - - future = asyncio.run(main()) - - assert not future.cancelled() - class TestS3Executor: @pytest.mark.parametrize("executor_class", [S3ThreadPoolExecutor, S3AioThreadPool]) @@ -626,32 +523,41 @@ def work(): def test_submit_to_loop_shutdown(self): # A callable that S3AioExecutor drops when its event loop shuts down - # cancels the future of the caller, so that wait() returns. + # is never run, and its future, whether submit() returned it or the + # caller passed it to submit_to(), is cancelled and settled, so that + # wait() returns, instead of left pending. events = [] started = threading.Event() - settled = threading.Event() + settled = threading.Semaphore(0) def work(): started.set() - settled.wait(5) + # Running until both pending futures are settled, so that their + # callables cannot start before the shutdown. + for _ in range(2): + assert settled.acquire(timeout=EVENT_TIMEOUT) events.append("finished") pending: Future[None] = Future() - pending.add_done_callback(lambda _: settled.set()) + pending.add_done_callback(lambda _: settled.release()) async def main(): + # The pending callables wait for the only thread of the executor. executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) running = executor.submit(work) + submitted = executor.submit(events.append, "submitted finished") + submitted.add_done_callback(lambda _: settled.release()) executor.submit_to(pending, events.append, "pending finished") - while not started.is_set(): - await asyncio.sleep(0.01) - return running + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + return running, submitted - running = asyncio.run(main()) + running, submitted = asyncio.run(main()) - _, not_done = wait([running, pending], timeout=5) + _, not_done = wait([running, submitted, pending], timeout=EVENT_TIMEOUT) assert not not_done assert events == ["finished"] + assert not running.cancelled() + assert submitted.cancelled() assert pending.cancelled() def test_submit_to_executor_shut_down(self): @@ -665,28 +571,8 @@ async def main(): executor._pool.shutdown() executor.submit_to(future, sum, [1, 2]) with pytest.raises(RuntimeError, match="after shutdown"): - await asyncio.wrap_future(future) + await asyncio.wait_for(asyncio.wrap_future(future), EVENT_TIMEOUT) asyncio.run(main()) assert not future.cancelled() - - def test_submit_to_cancelled_executor_shut_down(self): - # A future that the caller cancels before S3AioExecutor reports the - # error is marked as notified, so that wait() counts it as done. - future: Future[int] = Future() - - async def main(): - loop = asyncio.get_running_loop() - executor = S3AioExecutor(loop=loop) - # Fails the dispatch of the callable to the pool. - executor._pool.shutdown() - executor.submit_to(future, sum, [1, 2]) - assert future.cancel() - with ThreadPoolExecutor(max_workers=1) as waiter: - return await loop.run_in_executor(waiter, lambda: wait([future], timeout=5)) - - _, not_done = asyncio.run(main()) - - assert not not_done - assert future.cancelled() From 320d6cf372e352a96d0a5d5f5ce15e44c002739f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 23:37:25 +0900 Subject: [PATCH 08/19] Restore the async and executor checks that round one found uncovered (GH-1239) - test_fork: a thread of the parent holds the pool lock while the main thread forks, so that the child creates its own lock. - Restore test_submit_to_cancelled_executor_shut_down: a future that the caller cancelled before the dispatch failed must still be notified for wait(). - Check the code, message and HTTP status code of S3ClientError. - Compare the default max_workers of S3AioThreadPool with ThreadPoolExecutor's, which its docstring documents. - test_cancelled_read_frees_ranges cancels the read once the tasks of the other ranges have taken their results, and the max_workers copy test makes its first three part copies meet at a barrier. - Sign the presigned URLs without depending on botocore's signature version or addressing style, and build the test_get_version destinations in the test. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 48 +++++++++++++++---- tests/pyathena/filesystem/test_s3_core.py | 20 ++++++-- tests/pyathena/filesystem/test_s3_errors.py | 8 +++- tests/pyathena/filesystem/test_s3_executor.py | 47 +++++++++++++++++- 4 files changed, 104 insertions(+), 19 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index b965c9c1c..1817fcd27 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -223,12 +223,28 @@ async def test_copy_object_with_multipart_upload_max_workers(self): {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} ) ) - # Waits briefly for a fourth overlapping call, which only a copy that - # exceeds max_workers runs. - call, state = overlapping_calls(4, timeout=0.5) + condition = threading.Condition() + state = {"calls": 0, "active": 0, "peak": 0} + # The first three part copies run at once. + first_three = threading.Barrier(3, timeout=EVENT_TIMEOUT) def upload_part_copy(**kw): - call() + with condition: + state["calls"] += 1 + number = state["calls"] + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + condition.notify_all() + try: + if number <= 3: + first_three.wait() + # Waits briefly for a fourth overlapping call, which only a + # copy that exceeds max_workers runs. + with condition: + condition.wait_for(lambda: state["active"] > 3, timeout=0.5) + finally: + with condition: + state["active"] -= 1 return SimpleNamespace(etag='"e"', part_number=kw["part_number"]) sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) @@ -1361,10 +1377,10 @@ async def test_expand_path_glob_lists_stem_prefix(self): [ # GH-979: a version is downloaded to a local path named after its # key. - (lambda tmp_path: f"{tmp_path}/", "b"), + ("directory", "b"), # A Path destination is paired too: the version is downloaded to # the file, not into a directory of that name. - (lambda tmp_path: tmp_path / "out.bin", "out.bin"), + ("path", "out.bin"), ], ids=["directory", "path"], ) @@ -1374,7 +1390,8 @@ async def test_get_version(self, tmp_path, destination, expected): fs._sync_fs.isdir = mock.MagicMock(return_value=False) fs._get_file = mock.AsyncMock() - await fs._get("s3://bucket/b?versionId=v1", destination(tmp_path)) + lpath = f"{tmp_path}/" if destination == "directory" else tmp_path / "out.bin" + await fs._get("s3://bucket/b?versionId=v1", lpath) assert fs._get_file.await_args.args[:2] == ( "bucket/b?versionId=v1", f"{tmp_path.as_posix()}/{expected}", @@ -2560,15 +2577,26 @@ async def test_cancelled_read_frees_ranges(self, fs): held = threading.Event() release = threading.Event() + async def others_done(): + def running_reads(): + return [ + t + for t in asyncio.all_tasks() + if t.get_coro().__qualname__ == "S3AioThreadPool.run" and not t.done() + ] + + while len(running_reads()) > 1: + await asyncio.sleep(0) + def get_object(path, range_, **kwargs): # Unlike bytes, an array can be weakly referenced. data = array.array("B", bytes(range_[1] - range_[0])) read.append(weakref.ref(data)) all_read.wait() if range_[0] == 0: - # Held until the read is cancelled, after a pass of the event - # loop that takes the results of the other ranges. - asyncio.run_coroutine_threadsafe(asyncio.sleep(0), loop).result(EVENT_TIMEOUT) + # Held until the read is cancelled, once the tasks of the + # other ranges have taken their results. + asyncio.run_coroutine_threadsafe(others_done(), loop).result(EVENT_TIMEOUT) held.set() assert release.wait(EVENT_TIMEOUT) return data diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index 23e9046c8..ab02ddb3d 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -1701,12 +1701,22 @@ def test_generate_presigned_url(self, path, client_method, params, key, query): url = core.generate_presigned_url(path, client_method, expires_in=3600, **params) signed = time.time() parts = urlsplit(url) - assert (parts.netloc, parts.path) == ("bucket.s3.amazonaws.com", f"/{key}") + # The addressing style is botocore's. + assert parts.netloc.startswith("bucket.") or parts.path.startswith("/bucket/") + assert parts.path.endswith(f"/{key}") signed_query = parse_qs(parts.query) - # The signature fields of the client's signer, and the expiry. - expires = int(signed_query.pop("Expires")[0]) - assert signed + 3600 - 60 < expires <= signed + 3600 - del signed_query["AWSAccessKeyId"], signed_query["Signature"] + # The expiry, in the fields of the client's signature version. + if "X-Amz-Expires" in signed_query: + assert signed_query["X-Amz-Expires"] == ["3600"] + else: + expires = int(signed_query["Expires"][0]) + assert signed + 3600 - 60 < expires <= signed + 3600 + signature_fields = {"AWSAccessKeyId", "Signature", "Expires"} + signed_query = { + k: v + for k, v in signed_query.items() + if k not in signature_fields and not k.startswith("X-Amz-") + } # request_kwargs, such as RequestPayer, are not signed. assert signed_query == query diff --git a/tests/pyathena/filesystem/test_s3_errors.py b/tests/pyathena/filesystem/test_s3_errors.py index 2554d442a..68940cd81 100644 --- a/tests/pyathena/filesystem/test_s3_errors.py +++ b/tests/pyathena/filesystem/test_s3_errors.py @@ -42,7 +42,9 @@ class TestS3ClientError: ], ) def test_os_error_by_error_code(self, code, expected): - actual = S3ClientError(_client_error(code)).os_error + error = S3ClientError(_client_error(code)) + assert (error.code, error.message) == (code, "error message") + actual = error.os_error assert type(actual) is expected assert "error message" in str(actual) @@ -82,7 +84,9 @@ def test_os_error_by_error_code_errno(self, code, expected_errno): ], ) def test_os_error_by_http_status_code(self, status_code, expected_errno): - actual = S3ClientError(_client_error("UnknownCode", status_code=status_code)).os_error + error = S3ClientError(_client_error("UnknownCode", status_code=status_code)) + assert error.http_status_code == status_code + actual = error.os_error assert isinstance(actual, OSError) assert actual.errno == expected_errno diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 264aa856e..3b106d618 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -27,6 +27,11 @@ class TestS3AioThreadPool: + def test_default_max_workers(self): + # None takes the default of ThreadPoolExecutor, as documented. + with ThreadPoolExecutor() as executor: + assert S3AioThreadPool().max_workers == executor._max_workers + def test_run_beyond_default_executor(self): # GH-1159: the calls run in threads of the pool, so that max_workers, # not the default executor of the event loop, bounds them. @@ -188,10 +193,25 @@ def submit_and_notify(fn, *args, **kwargs): @pytest.mark.filterwarnings("ignore:.*fork\\(\\) may lead to deadlocks:DeprecationWarning") def test_fork(self): # A forked child, whose copy of the pool has no threads, creates new - # ones. + # ones, also if a thread of the parent held the lock of the pool. pool = S3AioThreadPool(1) assert pool.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 - pid = os.fork() + # No public call holds the lock across a fork; a thread of the parent + # holds it while the main thread forks. + held, release = threading.Event(), threading.Event() + + def hold_lock(): + with pool._lock: + held.set() + release.wait(EVENT_TIMEOUT) + + holder = threading.Thread(target=hold_lock) + holder.start() + try: + assert held.wait(EVENT_TIMEOUT) + pid = os.fork() + finally: + release.set() if pid == 0: # The child exits here, whatever happens, instead of running pytest. code = 1 @@ -212,6 +232,7 @@ def test_fork(self): if status == (0, 0): os.kill(pid, signal.SIGKILL) os.waitpid(pid, 0) + holder.join(EVENT_TIMEOUT) pool.shutdown() assert os.waitstatus_to_exitcode(status[1]) == 0 @@ -576,3 +597,25 @@ async def main(): asyncio.run(main()) assert not future.cancelled() + + def test_submit_to_cancelled_executor_shut_down(self): + # A future that the caller cancels before S3AioExecutor reports the + # error is marked as notified, so that wait() counts it as done. + future: Future[int] = Future() + + async def main(): + loop = asyncio.get_running_loop() + executor = S3AioExecutor(loop=loop) + # Fails the dispatch of the callable to the pool. + executor._pool.shutdown() + executor.submit_to(future, sum, [1, 2]) + assert future.cancel() + with ThreadPoolExecutor(max_workers=1) as waiter: + return await loop.run_in_executor( + waiter, lambda: wait([future], timeout=EVENT_TIMEOUT) + ) + + _, not_done = asyncio.run(main()) + + assert not not_done + assert future.cancelled() From 931ee351afaa65cf790848d3585d8fe3cf54ef8a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 23:56:54 +0900 Subject: [PATCH 09/19] Wake every overlapping call once the peak is reached (GH-1239) overlapping_calls() let the last call return before the others woke, which then waited out the whole timeout although the overlap had been reached; three executor and read tests waited 5 s each. Wait for the peak instead of the active calls. Also use wait_until() in test_cancelled_read_frees_ranges and check that its filter finds the held read. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 6 ++++-- tests/pyathena/util.py | 4 +++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 1817fcd27..de239e609 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -41,6 +41,7 @@ overlapping_calls, retried, stub_multipart_copy, + wait_until, ) @@ -2585,8 +2586,9 @@ def running_reads(): if t.get_coro().__qualname__ == "S3AioThreadPool.run" and not t.done() ] - while len(running_reads()) > 1: - await asyncio.sleep(0) + # At least the held read is found. + assert running_reads() + await wait_until(lambda: len(running_reads()) <= 1) def get_object(path, range_, **kwargs): # Unlike bytes, an array can be weakly referenced. diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 82e3d467c..62ddb6b08 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -222,7 +222,9 @@ def call(*args, **kwargs): state["active"] += 1 state["peak"] = max(state["peak"], state["active"]) condition.notify_all() - condition.wait_for(lambda: state["active"] >= calls, timeout=timeout) + # The peak, unlike the active calls, stays reached when the last + # call returns before the others wake up. + condition.wait_for(lambda: state["peak"] >= calls, timeout=timeout) with condition: state["active"] -= 1 From e635c90dc6f8af79eb87c98a8c71d1e164aaf835 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 00:45:40 +0900 Subject: [PATCH 10/19] Restore the async checks that the independent review found uncovered (GH-1239) - Restore TestS3AioThreadPool::test_shutdown: the pool refuses new calls after shutdown(wait=False) and runs the queued ones. - Restore the async test_multipart_copy_uses_creation_algorithm: the async copy keeps the checksums of the copied parts in the completion. - Send an explicit part checksum through S3Core.upload_part(). - The held range of test_cancelled_read_frees_ranges retries until the cancellation stops it, so that a read left uncancelled fails the test. - Presign with a 60 s expiry and a parameter that only PutObject accepts, so that the method and the expiry reach the signer. - Revert the overlapping_calls() change: callers that check an upper bound rely on the first calls waiting while later ones start. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 103 ++++++++++++++++-- tests/pyathena/filesystem/test_s3_core.py | 24 +++- tests/pyathena/filesystem/test_s3_executor.py | 15 +++ tests/pyathena/util.py | 4 +- 4 files changed, 126 insertions(+), 20 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index de239e609..64e67e393 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1551,6 +1551,81 @@ def test_internal_file_system_not_cached(self): for fs in (fs1, fs2, fs3): assert fs._sync_fs not in S3FileSystem._cache.values() + @pytest.mark.parametrize( + ("algorithm", "checksum_type"), + [(None, None), ("SHA256", None), ("CRC32", None), ("CRC32", "FULL_OBJECT")], + ) + @pytest.mark.asyncio + async def test_multipart_copy_uses_creation_algorithm(self, algorithm, checksum_type): + fs = AioS3FileSystem( + key="dummy", + secret="dummy", + region_name="us-east-1", + max_workers=1, + skip_instance_cache=True, + ) + block_size = 5 * 2**30 + size = 2 * block_size + checksum_kwargs = {"ChecksumAlgorithm": algorithm} if algorithm else {} + if checksum_type: + checksum_kwargs["ChecksumType"] = checksum_type + expected_parts = [] + with Stubber(fs.core.client) as stubber: + stubber.add_response( + "head_object", {"ContentLength": size}, {"Bucket": "bucket", "Key": "src"} + ) + stubber.add_response( + "create_multipart_upload", + {"Bucket": "bucket", "Key": "dst", "UploadId": "u", **checksum_kwargs}, + {"Bucket": "bucket", "Key": "dst", **checksum_kwargs}, + ) + for number, range_ in ( + (1, f"bytes=0-{block_size - 1}"), + (2, f"bytes={block_size}-{size - 1}"), + ): + result = {"ETag": f'"p{number}"', f"Checksum{algorithm or 'CRC32'}": "checksum"} + stubber.add_response( + "upload_part_copy", + {"CopyPartResult": result}, + { + "Bucket": "bucket", + "Key": "dst", + "UploadId": "u", + "PartNumber": number, + "CopySource": {"Bucket": "bucket", "Key": "src"}, + "CopySourceRange": range_, + }, + ) + expected_parts.append( + { + "ETag": result["ETag"], + "PartNumber": number, + **({f"Checksum{algorithm}": "checksum"} if algorithm else {}), + } + ) + stubber.add_response( + "complete_multipart_upload", + {"ETag": '"done"'}, + { + "Bucket": "bucket", + "Key": "dst", + "UploadId": "u", + "MultipartUpload": {"Parts": expected_parts}, + **({"ChecksumType": checksum_type} if checksum_type else {}), + }, + ) + kwargs = { + "source": S3Path("bucket", "src"), + "destination": S3Path("bucket", "dst"), + "block_size": block_size, + "MetadataDirective": "REPLACE", + "TaggingDirective": "REPLACE", + "AnnotationDirective": "EXCLUDE", + **checksum_kwargs, + } + await fs._copy_object_with_multipart_upload(**kwargs) + stubber.assert_no_pending_responses() + @pytest.fixture(scope="class") def fs(self): return AioS3FileSystem(connection=connect()) @@ -2569,14 +2644,16 @@ async def test_cancelled_creation_stops_retries(self, fs): @pytest.mark.asyncio async def test_cancelled_read_frees_ranges(self, fs): # GH-1186: a cancelled parallel read cancels the reads that have not - # finished and frees the data of the others, also while the - # cancellation is held, as a failed read does. + # finished, which stops their retries, and frees the data of the + # others, also while the cancellation is held, as a failed read does. block_size = 2**10 loop = asyncio.get_running_loop() read = [] all_read = threading.Barrier(4, timeout=EVENT_TIMEOUT) held = threading.Event() - release = threading.Event() + # The held range throttles and would retry after EVENT_TIMEOUT. + operation = ThrottledOperation(failures=1) + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) async def others_done(): def running_reads(): @@ -2596,11 +2673,12 @@ def get_object(path, range_, **kwargs): read.append(weakref.ref(data)) all_read.wait() if range_[0] == 0: - # Held until the read is cancelled, once the tasks of the + # Retrying when the read is cancelled, once the tasks of the # other ranges have taken their results. asyncio.run_coroutine_threadsafe(others_done(), loop).result(EVENT_TIMEOUT) held.set() - assert release.wait(EVENT_TIMEOUT) + with contextlib.suppress(Exception): + throttled() return data fs.core.get_object = mock.MagicMock(side_effect=get_object) @@ -2615,15 +2693,16 @@ def get_object(path, range_, **kwargs): f = await fs.open_async("s3://bucket/key", "rb", block_size=block_size, max_workers=4) task = asyncio.ensure_future(f.read()) - try: - assert await asyncio.to_thread(held.wait, EVENT_TIMEOUT) - task.cancel() - with pytest.raises(asyncio.CancelledError) as info_: - await task - finally: - release.set() + assert await asyncio.to_thread(held.wait, EVENT_TIMEOUT) + started = time.monotonic() + task.cancel() + with pytest.raises(asyncio.CancelledError) as info_: + await task await f.close() await asyncio.to_thread(f._pool.shutdown, True) + # The retries of the held range stopped. + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 for _ in range(3): await asyncio.sleep(0) assert len(read) == 4 diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index ab02ddb3d..709361968 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -1038,12 +1038,14 @@ def test_upload_part_uses_creation_algorithm_and_identity(self): "PartNumber": 1, "Body": b"data", "ChecksumAlgorithm": "SHA256", + "ChecksumSHA256": "c2hh", "SSECustomerAlgorithm": "AES256", }, ) with stubber: # The identity and the part number override the per-call fields, - # and the other parameters are sent. + # and the other parameters, such as the checksum of the part, are + # sent. part = core.upload_part( upload, 1, @@ -1053,6 +1055,7 @@ def test_upload_part_uses_creation_algorithm_and_identity(self): UploadId="other", PartNumber=99, ChecksumAlgorithm="CRC32", + ChecksumSHA256="c2hh", SSECustomerAlgorithm="AES256", ) stubber.assert_no_pending_responses() @@ -1691,14 +1694,21 @@ def test_replace_object_metadata_rejects_fields_of_the_copy(self, name): "other", {"versionId": ["v1"]}, ), - (S3Path("bucket", "key"), "put_object", {"Key": "other"}, "other", {}), + # A parameter that only PutObject accepts signs the method given. + ( + S3Path("bucket", "key"), + "put_object", + {"Key": "other", "ContentType": "text/csv"}, + "other", + {}, + ), ], ) def test_generate_presigned_url(self, path, client_method, params, key, query): # No request is sent, so the stubber has no responses to return. core, stubber = _make_core(request_kwargs={"RequestPayer": "requester"}) with stubber: - url = core.generate_presigned_url(path, client_method, expires_in=3600, **params) + url = core.generate_presigned_url(path, client_method, expires_in=60, **params) signed = time.time() parts = urlsplit(url) # The addressing style is botocore's. @@ -1707,10 +1717,14 @@ def test_generate_presigned_url(self, path, client_method, params, key, query): signed_query = parse_qs(parts.query) # The expiry, in the fields of the client's signature version. if "X-Amz-Expires" in signed_query: - assert signed_query["X-Amz-Expires"] == ["3600"] + assert signed_query["X-Amz-Expires"] == ["60"] else: expires = int(signed_query["Expires"][0]) - assert signed + 3600 - 60 < expires <= signed + 3600 + assert signed + 60 - 30 < expires <= signed + 60 + # SigV2 also signs the content type in the query. + assert signed_query.pop("content-type", [params.get("ContentType")]) == [ + params.get("ContentType") + ] signature_fields = {"AWSAccessKeyId", "Signature", "Expires"} signed_query = { k: v diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 3b106d618..660ffe896 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -189,6 +189,21 @@ def submit_and_notify(fn, *args, **kwargs): assert events == (["pending finished"] if shielded else []) + def test_shutdown(self): + # The calls submitted before the shutdown still run. + events = [] + release = threading.Event() + pool = S3AioThreadPool(1) + pool.submit(release.wait, EVENT_TIMEOUT) + pending = pool.submit(events.append, "pending finished") + pool.shutdown(wait=False) + with pytest.raises(RuntimeError, match="after shutdown"): + pool.submit(sum, [1, 2]) + release.set() + pending.result(timeout=EVENT_TIMEOUT) + + assert events == ["pending finished"] + @pytest.mark.skipif(not hasattr(os, "fork"), reason="os.fork() is not available.") @pytest.mark.filterwarnings("ignore:.*fork\\(\\) may lead to deadlocks:DeprecationWarning") def test_fork(self): diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 62ddb6b08..82e3d467c 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -222,9 +222,7 @@ def call(*args, **kwargs): state["active"] += 1 state["peak"] = max(state["peak"], state["active"]) condition.notify_all() - # The peak, unlike the active calls, stays reached when the last - # call returns before the others wake up. - condition.wait_for(lambda: state["peak"] >= calls, timeout=timeout) + condition.wait_for(lambda: state["active"] >= calls, timeout=timeout) with condition: state["active"] -= 1 From 117e12f1e5beca99a5437365f7696a188b58c28e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 01:05:54 +0900 Subject: [PATCH 11/19] Bound and tighten the repaired async and executor tests (GH-1239) - test_cancelled_read_frees_ranges creates the data of the held range after its stopped retry, whose traceback would otherwise keep it, and bounds its waits for the cancelled read and the pool shutdown. - test_shutdown checks that shutdown(wait=False) does not wait for the running call, and no longer matches the error message. - test_generate_presigned_url checks that the content type is signed, in the query with SigV2 and as a signed header with SigV4. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 18 +++++++++++------- tests/pyathena/filesystem/test_s3_core.py | 14 ++++++++++---- tests/pyathena/filesystem/test_s3_executor.py | 4 +++- 3 files changed, 24 insertions(+), 12 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 64e67e393..9f61c7aef 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2668,17 +2668,21 @@ def running_reads(): await wait_until(lambda: len(running_reads()) <= 1) def get_object(path, range_, **kwargs): - # Unlike bytes, an array can be weakly referenced. - data = array.array("B", bytes(range_[1] - range_[0])) - read.append(weakref.ref(data)) - all_read.wait() if range_[0] == 0: + all_read.wait() # Retrying when the read is cancelled, once the tasks of the - # other ranges have taken their results. + # other ranges have taken their results. Its data is created + # afterwards, so that the traceback of the stopped retry does + # not keep it. asyncio.run_coroutine_threadsafe(others_done(), loop).result(EVENT_TIMEOUT) held.set() with contextlib.suppress(Exception): throttled() + # Unlike bytes, an array can be weakly referenced. + data = array.array("B", bytes(range_[1] - range_[0])) + read.append(weakref.ref(data)) + if range_[0] != 0: + all_read.wait() return data fs.core.get_object = mock.MagicMock(side_effect=get_object) @@ -2697,9 +2701,9 @@ def get_object(path, range_, **kwargs): started = time.monotonic() task.cancel() with pytest.raises(asyncio.CancelledError) as info_: - await task + await asyncio.wait_for(task, EVENT_TIMEOUT) await f.close() - await asyncio.to_thread(f._pool.shutdown, True) + await asyncio.wait_for(asyncio.to_thread(f._pool.shutdown, True), EVENT_TIMEOUT) # The retries of the held range stopped. assert time.monotonic() - started < EVENT_TIMEOUT assert operation.calls == 1 diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index 709361968..7780322a2 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -1721,10 +1721,16 @@ def test_generate_presigned_url(self, path, client_method, params, key, query): else: expires = int(signed_query["Expires"][0]) assert signed + 60 - 30 < expires <= signed + 60 - # SigV2 also signs the content type in the query. - assert signed_query.pop("content-type", [params.get("ContentType")]) == [ - params.get("ContentType") - ] + # The content type is signed: in the query with SigV2, as a header + # with SigV4. + content_type = params.get("ContentType") + if "Signature" in signed_query: + assert signed_query.pop("content-type", None) == ( + [content_type] if content_type else None + ) + else: + signed_headers = signed_query["X-Amz-SignedHeaders"][0].split(";") + assert ("content-type" in signed_headers) is (content_type is not None) signature_fields = {"AWSAccessKeyId", "Signature", "Expires"} signed_query = { k: v diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 660ffe896..e93de804f 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -197,7 +197,9 @@ def test_shutdown(self): pool.submit(release.wait, EVENT_TIMEOUT) pending = pool.submit(events.append, "pending finished") pool.shutdown(wait=False) - with pytest.raises(RuntimeError, match="after shutdown"): + # The shutdown did not wait for the running call. + assert not pending.done() + with pytest.raises(RuntimeError): pool.submit(sum, [1, 2]) release.set() pending.result(timeout=EVENT_TIMEOUT) From 5dc59fc3072de94c3f9819c48da8e6bfef6fba1d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 01:24:47 +0900 Subject: [PATCH 12/19] Break the reference cycle of the stopped retry in the cancelled read test (GH-1239) The error of a retry that the cancellation stops is kept in a cycle through the frames of its traceback, which kept the frame of the held range and its data until the cyclic collector ran. Clear those frames, as S3Core.call() does. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 9f61c7aef..be122c0c5 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -7,6 +7,7 @@ import tempfile import threading import time +import traceback import uuid import weakref from concurrent.futures import Future, ThreadPoolExecutor, wait @@ -2671,13 +2672,16 @@ def get_object(path, range_, **kwargs): if range_[0] == 0: all_read.wait() # Retrying when the read is cancelled, once the tasks of the - # other ranges have taken their results. Its data is created - # afterwards, so that the traceback of the stopped retry does - # not keep it. + # other ranges have taken their results. asyncio.run_coroutine_threadsafe(others_done(), loop).result(EVENT_TIMEOUT) held.set() - with contextlib.suppress(Exception): + try: throttled() + except Exception as e: + # The stopped retry keeps its error in a reference cycle + # through the frames of its traceback, which would keep + # this frame and the data, as S3Core.call() avoids. + traceback.clear_frames(e.__traceback__) # Unlike bytes, an array can be weakly referenced. data = array.array("B", bytes(range_[1] - range_[0])) read.append(weakref.ref(data)) From eb923cf7d39909fa17622a237b8b03853f1cada3 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 09:29:12 +0900 Subject: [PATCH 13/19] Bound the one-worker copy tests and let overlapping_calls see an exceeded bound (GH-1239) - Await the copies of the tests that run with one worker within EVENT_TIMEOUT, so that a copy that never gets a worker fails the test instead of hanging it. - overlapping_calls() waits for the peak instead of the active calls, and then up to a grace period for one more call, which only a caller that runs more calls at once makes; a pool one worker too large now fails test_run_beyond_default_executor and test_read_parallel. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 67 +++++++++++++--------- tests/pyathena/util.py | 15 +++-- 2 files changed, 51 insertions(+), 31 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index be122c0c5..b8e8049c7 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -71,11 +71,14 @@ async def _multipart_copy(fs=None, **kwargs): with Stubber(fs._sync_fs._client) as stubber: stub_multipart_copy(stubber, **kwargs) try: - await fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), - S3Path("bucket", "dst"), - block_size=MULTIPART_COPY_BLOCK_SIZE, - **MULTIPART_COPY_KWARGS, + await asyncio.wait_for( + fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + block_size=MULTIPART_COPY_BLOCK_SIZE, + **MULTIPART_COPY_KWARGS, + ), + EVENT_TIMEOUT, ) finally: stubber.assert_no_pending_responses() @@ -199,12 +202,15 @@ def upload_part_copy(**kw): ) with pytest.raises(OSError, match="part failed"): - await fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), - S3Path("bucket", "dst"), - MetadataDirective="REPLACE", - TaggingDirective="REPLACE", - AnnotationDirective="EXCLUDE", + await asyncio.wait_for( + fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ), + EVENT_TIMEOUT, ) # Part 3 waits for a worker and is not started after the failure. @@ -707,8 +713,9 @@ def upload_part_copy(part_number, **kwargs): sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) started = time.monotonic() with pytest.raises(PermissionError, match="Access Denied"): - await fs._copy_object_with_multipart_upload( - plan.source, plan.destination, max_workers=2 + await asyncio.wait_for( + fs._copy_object_with_multipart_upload(plan.source, plan.destination, max_workers=2), + EVENT_TIMEOUT, ) assert time.monotonic() - started < EVENT_TIMEOUT assert operation.calls == 1 @@ -788,7 +795,9 @@ def copy_object_annotation(name, *args, **kwargs): sync_fs.core.copy_object_annotation = mock.MagicMock(side_effect=copy_object_annotation) started = time.monotonic() with pytest.raises(PermissionError, match="Access Denied"): - await fs._copy_object_with_multipart_upload(plan.source, plan.destination) + await asyncio.wait_for( + fs._copy_object_with_multipart_upload(plan.source, plan.destination), EVENT_TIMEOUT + ) assert time.monotonic() - started < EVENT_TIMEOUT assert operation.calls == 1 @@ -828,8 +837,9 @@ def copy_object_annotation(name, *args, **kwargs): sync_fs.core.copy_object_annotation = mock.MagicMock(side_effect=copy_object_annotation) with pytest.raises(PermissionError, match="Access Denied"): - await fs._copy_object_with_multipart_upload( - plan.source, plan.destination, max_workers=3 + await asyncio.wait_for( + fs._copy_object_with_multipart_upload(plan.source, plan.destination, max_workers=3), + EVENT_TIMEOUT, ) sync_fs.core.copy_object_annotation.assert_called_once() @@ -1500,16 +1510,19 @@ def upload_part_copy(**kw): "AnnotationDirective": "EXCLUDE", } - await fs._cp_file( - "s3://bucket/src", - "s3://bucket/dst", - block_size=S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE // 2, - max_workers=1, - RequestPayer="requester", - ContentType="text/csv", - # Copy without reading the metadata, tags and annotations of the - # source (GH-973). - **directives, + await asyncio.wait_for( + fs._cp_file( + "s3://bucket/src", + "s3://bucket/dst", + block_size=S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE // 2, + max_workers=1, + RequestPayer="requester", + ContentType="text/csv", + # Copy without reading the metadata, tags and annotations of the + # source (GH-973). + **directives, + ), + EVENT_TIMEOUT, ) if size <= S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE: @@ -1624,7 +1637,7 @@ async def test_multipart_copy_uses_creation_algorithm(self, algorithm, checksum_ "AnnotationDirective": "EXCLUDE", **checksum_kwargs, } - await fs._copy_object_with_multipart_upload(**kwargs) + await asyncio.wait_for(fs._copy_object_with_multipart_upload(**kwargs), EVENT_TIMEOUT) stubber.assert_no_pending_responses() @pytest.fixture(scope="class") diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 82e3d467c..ce519fe7d 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -201,14 +201,18 @@ async def wait_until(predicate, timeout=EVENT_TIMEOUT): await asyncio.sleep(0.01) -def overlapping_calls(calls, timeout=5): +def overlapping_calls(calls, timeout=5, grace=0.5): """Return a callable that records the peak number of its overlapping calls. - Each call waits until ``calls`` calls overlap, or until the timeout. + Each call waits until ``calls`` calls overlap, or until the timeout. Once + they do, it waits up to ``grace`` seconds more for another overlapping + call, which only a caller that runs more than ``calls`` at once makes, so + that the peak also shows an exceeded bound. Args: calls: The number of calls to wait for. - timeout: The longest wait of a call, in seconds. + timeout: The longest wait of a call for the overlap, in seconds. + grace: The longest wait of a call for one more, in seconds. Returns: The callable, which accepts and ignores any arguments, and a @@ -222,7 +226,10 @@ def call(*args, **kwargs): state["active"] += 1 state["peak"] = max(state["peak"], state["active"]) condition.notify_all() - condition.wait_for(lambda: state["active"] >= calls, timeout=timeout) + # The peak, unlike the active calls, stays reached when the last + # call returns before the others wake up. + if condition.wait_for(lambda: state["peak"] >= calls, timeout=timeout): + condition.wait_for(lambda: state["peak"] > calls, timeout=grace) with condition: state["active"] -= 1 From 7a5f120e1d992912c7945087f90eb6cb6ac8aeeb Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 09:36:34 +0900 Subject: [PATCH 14/19] Check that the async delegations pass their arguments through (GH-1239) The async methods that run a method of the synchronous filesystem had no test of the arguments they pass: dropping the range of _cat_file, the mode of _pipe_file or the maxdepth of _mv failed nothing, because the sync tests cannot see the async forwarding. One parametrized test calls each of them and its synchronous method with the same arguments on filesystems that answer from the same keys, and compares the requests sent, the results and the stored keys. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 121 +++++++++++++++++++++ 1 file changed, 121 insertions(+) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index b8e8049c7..4eee31833 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -3,6 +3,7 @@ import contextlib import gc import gzip +import io import os import tempfile import threading @@ -28,6 +29,7 @@ S3MultipartUploadPart, S3Object, S3ObjectType, + S3ObjectVersion, ) from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_writer import S3MultipartStateError @@ -335,6 +337,125 @@ async def test_blocking_calls_bypass_default_executor(self): assert info == {"name": "bucket/key"} + @staticmethod + def _serve_keys(fs, keys): + # Answer the S3 requests of the filesystem from a set of keys of + # "bucket", which is returned, and record them by operation name. + store = set(keys) + requests = [] + + def call(method, **request): + name = method if isinstance(method, str) else method.__name__ + requests.append((name, request)) + key = request.get("Key") + if name == "head_object": + if key not in store: + raise FileNotFoundError(key) + return {"ContentLength": 2, "ETag": '"e"'} + if name == "get_object": + return {"Body": io.BytesIO(b"aa"), "ContentLength": 2, "ETag": '"e"'} + if name == "list_objects_v2": + prefix, delimiter = request.get("Prefix", ""), request.get("Delimiter") + listed = sorted(k for k in store if k.startswith(prefix)) + directories = sorted( + { + prefix + k[len(prefix) :].split(delimiter)[0] + delimiter + for k in listed + if delimiter and delimiter in k[len(prefix) :] + } + ) + files = [k for k in listed if not delimiter or delimiter not in k[len(prefix) :]] + return { + "Contents": [{"Key": k, "Size": 2, "ETag": '"e"'} for k in files], + "CommonPrefixes": [{"Prefix": d} for d in directories], + "KeyCount": len(files) + len(directories), + } + if name in {"put_object", "copy_object"}: + store.add(key) + return {"ETag": '"e"'} + if name == "delete_object": + store.discard(key) + if name == "delete_objects": + deleted = [o["Key"] for o in request["Delete"]["Objects"]] + store.difference_update(deleted) + return {"Deleted": [{"Key": k} for k in deleted]} + if name == "get_object_tagging": + return {"TagSet": [{"Key": "a", "Value": "1"}]} + if name == "list_object_versions": + return {"Versions": [{"Key": "key", "VersionId": "v1", "IsLatest": True}]} + return {} + + fs._sync_fs._call = fs._sync_fs._core.call = call + return store, requests + + @pytest.mark.parametrize( + ("method", "sync_method", "args", "kwargs"), + [ + ("_info", "info", ("s3://bucket/key?versionId=v1",), {}), + ("_ls", "ls", ("s3://bucket/dir",), {"detail": True}), + ("_cat_file", "cat_file", ("s3://bucket/key",), {"start": 1, "end": 2}), + ("_exists", "exists", ("s3://bucket/key?versionId=v1",), {}), + ("_rm_file", "rm_file", ("s3://bucket/key?versionId=v1",), {}), + ( + "_pipe_file", + "pipe_file", + ("s3://bucket/new", b"a"), + {"mode": "create", "ContentType": "text/csv"}, + ), + ("_put_file", "put_file", ("local-source", "s3://bucket/new"), {"mode": "create"}), + ("_get_file", "get_file", ("s3://bucket/key?versionId=v1", "local-destination"), {}), + ("_touch", "touch", ("s3://bucket/new",), {"ContentType": "text/csv"}), + ("_find", "find", ("s3://bucket/dir",), {"maxdepth": 1, "withdirs": True}), + ( + "_mv", + "mv", + ("s3://bucket/dir", "s3://bucket/out"), + {"recursive": True, "maxdepth": 1}, + ), + ("metadata", "metadata", ("s3://bucket/key?versionId=v1",), {}), + ("get_tags", "get_tags", ("s3://bucket/key?versionId=v1",), {}), + ("put_tags", "put_tags", ("s3://bucket/key", {"b": "2"}), {"mode": "m"}), + ("chmod", "chmod", ("s3://bucket/dir", "private"), {"recursive": True}), + ("object_version_info", "object_version_info", ("s3://bucket/key",), {}), + ], + ) + @pytest.mark.asyncio + async def test_delegations_send_the_requests_of_the_sync_methods( + self, tmp_path, method, sync_method, args, kwargs + ): + # The methods that run the synchronous filesystem's methods pass their + # arguments through: each sends the requests, and returns the result, + # of the synchronous method called with the same arguments. + keys = {"key", "dir/a", "dir/sub/b"} + source = tmp_path / "source" + source.write_bytes(b"a") + + def comparable(value): + # The listed versions have no equality of their own. + if isinstance(value, list): + return [comparable(v) for v in value] + if isinstance(value, S3ObjectVersion): + return vars(value) + return value + + outcomes = [] + for run, name in (("sync", sync_method), ("async", method)): + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + store, requests = self._serve_keys(fs, keys) + local = {"local-source": str(source), "local-destination": str(tmp_path / run)} + call_args = [local.get(a, a) if isinstance(a, str) else a for a in args] + if run == "sync": + result = getattr(fs._sync_fs, name)(*call_args, **kwargs) + else: + result = getattr(fs, name)(*call_args, **kwargs) + if asyncio.iscoroutine(result): + result = await asyncio.wait_for(result, EVENT_TIMEOUT) + outcomes.append((comparable(result), sorted(map(repr, requests)), sorted(store))) + assert outcomes[0][1], "no request was sent" + assert outcomes[1] == outcomes[0] + @pytest.mark.asyncio async def test_max_workers_none(self): # None takes the default of ThreadPoolExecutor for the threads of the From 6e22345149c92c74c059d820957846710639f64b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 09:50:58 +0900 Subject: [PATCH 15/19] Fail the bounded awaits without waiting for the cancellation (GH-1239) Round one of the review of the new checks found: - asyncio.wait_for() waits for the cancelled copy, whose cleanup a stuck worker can hold; await_bounded() fails the test at the bound instead. The elapsed-time checks that it makes redundant are gone. - The delegation fake answered copies of deleted keys and was used from several threads without a lock; it now rejects a copy of a missing source and serializes its requests. Each case passes a request parameter, so that dropping **kwargs fails it, and object_version_info is called with delete_markers. - overlapping_calls() makes only the calls of the first overlap wait for one more. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 109 ++++++++++++++------- tests/pyathena/util.py | 39 ++++++-- 2 files changed, 106 insertions(+), 42 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 4eee31833..17b12e8ef 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -41,6 +41,7 @@ MULTIPART_COPY_BLOCK_SIZE, MULTIPART_COPY_KWARGS, ThrottledOperation, + await_bounded, overlapping_calls, retried, stub_multipart_copy, @@ -73,14 +74,13 @@ async def _multipart_copy(fs=None, **kwargs): with Stubber(fs._sync_fs._client) as stubber: stub_multipart_copy(stubber, **kwargs) try: - await asyncio.wait_for( + await await_bounded( fs._copy_object_with_multipart_upload( S3Path("bucket", "src"), S3Path("bucket", "dst"), block_size=MULTIPART_COPY_BLOCK_SIZE, **MULTIPART_COPY_KWARGS, - ), - EVENT_TIMEOUT, + ) ) finally: stubber.assert_no_pending_responses() @@ -204,15 +204,14 @@ def upload_part_copy(**kw): ) with pytest.raises(OSError, match="part failed"): - await asyncio.wait_for( + await await_bounded( fs._copy_object_with_multipart_upload( S3Path("bucket", "src"), S3Path("bucket", "dst"), MetadataDirective="REPLACE", TaggingDirective="REPLACE", AnnotationDirective="EXCLUDE", - ), - EVENT_TIMEOUT, + ) ) # Part 3 waits for a worker and is not started after the failure. @@ -340,12 +339,17 @@ async def test_blocking_calls_bypass_default_executor(self): @staticmethod def _serve_keys(fs, keys): # Answer the S3 requests of the filesystem from a set of keys of - # "bucket", which is returned, and record them by operation name. + # "bucket", which is returned, and record them by operation name. The + # requests may come from several threads at once. store = set(keys) requests = [] + lock = threading.Lock() def call(method, **request): - name = method if isinstance(method, str) else method.__name__ + with lock: + return answer(method if isinstance(method, str) else method.__name__, request) + + def answer(name, request): requests.append((name, request)) key = request.get("Key") if name == "head_object": @@ -356,7 +360,7 @@ def call(method, **request): return {"Body": io.BytesIO(b"aa"), "ContentLength": 2, "ETag": '"e"'} if name == "list_objects_v2": prefix, delimiter = request.get("Prefix", ""), request.get("Delimiter") - listed = sorted(k for k in store if k.startswith(prefix)) + listed = sorted(k for k in list(store) if k.startswith(prefix)) directories = sorted( { prefix + k[len(prefix) :].split(delimiter)[0] + delimiter @@ -370,6 +374,8 @@ def call(method, **request): "CommonPrefixes": [{"Prefix": d} for d in directories], "KeyCount": len(files) + len(directories), } + if name == "copy_object" and request["CopySource"]["Key"] not in store: + raise FileNotFoundError(request["CopySource"]["Key"]) if name in {"put_object", "copy_object"}: store.add(key) return {"ETag": '"e"'} @@ -382,7 +388,10 @@ def call(method, **request): if name == "get_object_tagging": return {"TagSet": [{"Key": "a", "Value": "1"}]} if name == "list_object_versions": - return {"Versions": [{"Key": "key", "VersionId": "v1", "IsLatest": True}]} + return { + "Versions": [{"Key": "key", "VersionId": "v1", "IsLatest": False}], + "DeleteMarkers": [{"Key": "key", "VersionId": "m1", "IsLatest": True}], + } return {} fs._sync_fs._call = fs._sync_fs._core.call = call @@ -391,10 +400,25 @@ def call(method, **request): @pytest.mark.parametrize( ("method", "sync_method", "args", "kwargs"), [ - ("_info", "info", ("s3://bucket/key?versionId=v1",), {}), + ( + "_info", + "info", + ("s3://bucket/key?versionId=v1",), + {"ExpectedBucketOwner": "111122223333"}, + ), ("_ls", "ls", ("s3://bucket/dir",), {"detail": True}), - ("_cat_file", "cat_file", ("s3://bucket/key",), {"start": 1, "end": 2}), - ("_exists", "exists", ("s3://bucket/key?versionId=v1",), {}), + ( + "_cat_file", + "cat_file", + ("s3://bucket/key",), + {"start": 1, "end": 2, "ExpectedBucketOwner": "111122223333"}, + ), + ( + "_exists", + "exists", + ("s3://bucket/key?versionId=v1",), + {"ExpectedBucketOwner": "111122223333"}, + ), ("_rm_file", "rm_file", ("s3://bucket/key?versionId=v1",), {}), ( "_pipe_file", @@ -402,21 +426,41 @@ def call(method, **request): ("s3://bucket/new", b"a"), {"mode": "create", "ContentType": "text/csv"}, ), - ("_put_file", "put_file", ("local-source", "s3://bucket/new"), {"mode": "create"}), - ("_get_file", "get_file", ("s3://bucket/key?versionId=v1", "local-destination"), {}), + ( + "_put_file", + "put_file", + ("local-source", "s3://bucket/new"), + {"mode": "create", "ContentType": "text/csv"}, + ), + ( + "_get_file", + "get_file", + ("s3://bucket/key?versionId=v1", "local-destination"), + {"ExpectedBucketOwner": "111122223333"}, + ), ("_touch", "touch", ("s3://bucket/new",), {"ContentType": "text/csv"}), ("_find", "find", ("s3://bucket/dir",), {"maxdepth": 1, "withdirs": True}), ( "_mv", "mv", ("s3://bucket/dir", "s3://bucket/out"), - {"recursive": True, "maxdepth": 1}, + {"recursive": True, "maxdepth": 1, "MetadataDirective": "REPLACE"}, + ), + ( + "metadata", + "metadata", + ("s3://bucket/key?versionId=v1",), + {"ExpectedBucketOwner": "111122223333"}, ), - ("metadata", "metadata", ("s3://bucket/key?versionId=v1",), {}), ("get_tags", "get_tags", ("s3://bucket/key?versionId=v1",), {}), ("put_tags", "put_tags", ("s3://bucket/key", {"b": "2"}), {"mode": "m"}), ("chmod", "chmod", ("s3://bucket/dir", "private"), {"recursive": True}), - ("object_version_info", "object_version_info", ("s3://bucket/key",), {}), + ( + "object_version_info", + "object_version_info", + ("s3://bucket/key",), + {"delete_markers": True}, + ), ], ) @pytest.mark.asyncio @@ -451,7 +495,7 @@ def comparable(value): else: result = getattr(fs, name)(*call_args, **kwargs) if asyncio.iscoroutine(result): - result = await asyncio.wait_for(result, EVENT_TIMEOUT) + result = await await_bounded(result) outcomes.append((comparable(result), sorted(map(repr, requests)), sorted(store))) assert outcomes[0][1], "no request was sent" assert outcomes[1] == outcomes[0] @@ -832,13 +876,11 @@ def upload_part_copy(part_number, **kwargs): return SimpleNamespace(etag='"e"', part_number=part_number) sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) - started = time.monotonic() + # The retries stop before the bound of the wait. with pytest.raises(PermissionError, match="Access Denied"): - await asyncio.wait_for( - fs._copy_object_with_multipart_upload(plan.source, plan.destination, max_workers=2), - EVENT_TIMEOUT, + await await_bounded( + fs._copy_object_with_multipart_upload(plan.source, plan.destination, max_workers=2) ) - assert time.monotonic() - started < EVENT_TIMEOUT assert operation.calls == 1 sync_fs._abort_multipart_upload.assert_called_once() @@ -914,12 +956,11 @@ def copy_object_annotation(name, *args, **kwargs): return throttled() sync_fs.core.copy_object_annotation = mock.MagicMock(side_effect=copy_object_annotation) - started = time.monotonic() + # The retries stop before the bound of the wait. with pytest.raises(PermissionError, match="Access Denied"): - await asyncio.wait_for( - fs._copy_object_with_multipart_upload(plan.source, plan.destination), EVENT_TIMEOUT + await await_bounded( + fs._copy_object_with_multipart_upload(plan.source, plan.destination) ) - assert time.monotonic() - started < EVENT_TIMEOUT assert operation.calls == 1 @pytest.mark.asyncio @@ -958,9 +999,8 @@ def copy_object_annotation(name, *args, **kwargs): sync_fs.core.copy_object_annotation = mock.MagicMock(side_effect=copy_object_annotation) with pytest.raises(PermissionError, match="Access Denied"): - await asyncio.wait_for( - fs._copy_object_with_multipart_upload(plan.source, plan.destination, max_workers=3), - EVENT_TIMEOUT, + await await_bounded( + fs._copy_object_with_multipart_upload(plan.source, plan.destination, max_workers=3) ) sync_fs.core.copy_object_annotation.assert_called_once() @@ -1631,7 +1671,7 @@ def upload_part_copy(**kw): "AnnotationDirective": "EXCLUDE", } - await asyncio.wait_for( + await await_bounded( fs._cp_file( "s3://bucket/src", "s3://bucket/dst", @@ -1642,8 +1682,7 @@ def upload_part_copy(**kw): # Copy without reading the metadata, tags and annotations of the # source (GH-973). **directives, - ), - EVENT_TIMEOUT, + ) ) if size <= S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE: @@ -1758,7 +1797,7 @@ async def test_multipart_copy_uses_creation_algorithm(self, algorithm, checksum_ "AnnotationDirective": "EXCLUDE", **checksum_kwargs, } - await asyncio.wait_for(fs._copy_object_with_multipart_upload(**kwargs), EVENT_TIMEOUT) + await await_bounded(fs._copy_object_with_multipart_upload(**kwargs)) stubber.assert_no_pending_responses() @pytest.fixture(scope="class") diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index ce519fe7d..97b8b0e9c 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -201,25 +201,47 @@ async def wait_until(predicate, timeout=EVENT_TIMEOUT): await asyncio.sleep(0.01) +async def await_bounded(awaitable, timeout=EVENT_TIMEOUT): + """Await an awaitable, failing if it is not done within the timeout. + + Unlike ``asyncio.wait_for``, it does not wait for the cancellation of an + awaitable that the timeout interrupts, which a regression could also + hold, so that the test fails instead of hanging. + + Args: + awaitable: The coroutine or future to await. + timeout: The longest wait, in seconds. + + Returns: + The result of the awaitable. + """ + task = asyncio.ensure_future(awaitable) + done, _ = await asyncio.wait({task}, timeout=timeout) + assert done, f"not done within {timeout} s" + return task.result() + + def overlapping_calls(calls, timeout=5, grace=0.5): """Return a callable that records the peak number of its overlapping calls. - Each call waits until ``calls`` calls overlap, or until the timeout. Once - they do, it waits up to ``grace`` seconds more for another overlapping - call, which only a caller that runs more than ``calls`` at once makes, so - that the peak also shows an exceeded bound. + Until ``calls`` calls have overlapped, each call waits for that, or until + the timeout. The calls of that first overlap then wait up to ``grace`` + seconds for another overlapping call, which only a caller that runs more + than ``calls`` at once makes, so that the peak also shows an exceeded + bound. Later calls do not wait. Args: calls: The number of calls to wait for. timeout: The longest wait of a call for the overlap, in seconds. - grace: The longest wait of a call for one more, in seconds. + grace: The longest wait of a call of the first overlap for one more, + in seconds. Returns: The callable, which accepts and ignores any arguments, and a dictionary whose ``peak`` item is the peak number. """ condition = threading.Condition() - state = {"active": 0, "peak": 0} + state = {"active": 0, "peak": 0, "released": False} def call(*args, **kwargs): with condition: @@ -228,8 +250,11 @@ def call(*args, **kwargs): condition.notify_all() # The peak, unlike the active calls, stays reached when the last # call returns before the others wake up. - if condition.wait_for(lambda: state["peak"] >= calls, timeout=timeout): + if not state["released"] and condition.wait_for( + lambda: state["peak"] >= calls, timeout=timeout + ): condition.wait_for(lambda: state["peak"] > calls, timeout=grace) + state["released"] = True with condition: state["active"] -= 1 From 5dc3ca6082d0c633651fc448bdf9ad22a7a1e6d9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 10:18:57 +0900 Subject: [PATCH 16/19] Compare the async delegations by everything they return or write (GH-1239) The independent review found forwarding that the comparison missed: - a wrong download destination, or a dropped progress callback, of _get_file and _put_file: the downloaded files and the callback updates are compared too; - a dropped truncate of _touch: a case touches an existing key with truncate=False, and a raised error is compared by its type; - metadata() returning a plain mapping: S3Metadata and the listed versions compare by their public properties. A blocked thread can also hang a test after await_bounded() fails it, while pytest-asyncio cancels the task in the teardown. A per-test faulthandler watchdog ends the process after 300 s instead, which fails the test. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 83 ++++++++++++++++++---- 1 file changed, 70 insertions(+), 13 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 17b12e8ef..5a501be96 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1,6 +1,7 @@ import array import asyncio import contextlib +import faulthandler import gc import gzip import io @@ -19,12 +20,14 @@ import fsspec import pytest from botocore.stub import Stubber +from fsspec.callbacks import Callback from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile from pyathena.filesystem.s3_core import S3Core, S3MultipartCopyPlan from pyathena.filesystem.s3_executor import S3AioThreadPool, S3Executor from pyathena.filesystem.s3_object import ( + S3Metadata, S3MultipartUpload, S3MultipartUploadPart, S3Object, @@ -48,6 +51,37 @@ wait_until, ) +# A regression that leaves a thread blocked can hang a test in its body or in +# its teardown, where no timeout of the test can end it; the watchdog dumps the +# stacks and ends the process instead, which fails the test. +WATCHDOG_TIMEOUT = 300 + + +@pytest.fixture(autouse=True) +def watchdog(): + # Rearmed for each test, so that it also covers the teardown. + faulthandler.dump_traceback_later(WATCHDOG_TIMEOUT, exit=True) + + +@pytest.fixture(autouse=True, scope="module") +def watchdog_cancellation(): + yield + faulthandler.cancel_dump_traceback_later() + + +class _RecordingCallback(Callback): + """A progress callback that records its updates.""" + + def __init__(self): + super().__init__() + self.events = [] + + def set_size(self, size): + self.events.append(("size", size)) + + def relative_update(self, inc=1): + self.events.append(("update", inc)) + @pytest.fixture(scope="class") def register_async_filesystem(): @@ -430,15 +464,17 @@ def answer(name, request): "_put_file", "put_file", ("local-source", "s3://bucket/new"), - {"mode": "create", "ContentType": "text/csv"}, + {"mode": "create", "callback": "recording", "ContentType": "text/csv"}, ), ( "_get_file", "get_file", ("s3://bucket/key?versionId=v1", "local-destination"), - {"ExpectedBucketOwner": "111122223333"}, + {"callback": "recording", "ExpectedBucketOwner": "111122223333"}, ), ("_touch", "touch", ("s3://bucket/new",), {"ContentType": "text/csv"}), + # An existing object is not replaced. + ("_touch", "touch", ("s3://bucket/key",), {"truncate": False}), ("_find", "find", ("s3://bucket/dir",), {"maxdepth": 1, "withdirs": True}), ( "_mv", @@ -475,11 +511,14 @@ async def test_delegations_send_the_requests_of_the_sync_methods( source.write_bytes(b"a") def comparable(value): - # The listed versions have no equality of their own. + # The metadata and the listed versions compare by their public + # properties. if isinstance(value, list): return [comparable(v) for v in value] - if isinstance(value, S3ObjectVersion): - return vars(value) + if isinstance(value, (S3Metadata, S3ObjectVersion)): + cls = type(value) + properties = [n for n in dir(cls) if isinstance(getattr(cls, n), property)] + return cls, {n: getattr(value, n) for n in properties} return value outcomes = [] @@ -488,15 +527,33 @@ def comparable(value): key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True ) store, requests = self._serve_keys(fs, keys) - local = {"local-source": str(source), "local-destination": str(tmp_path / run)} + (tmp_path / run).mkdir() + local = { + "local-source": str(source), + "local-destination": str(tmp_path / run / "downloaded"), + } call_args = [local.get(a, a) if isinstance(a, str) else a for a in args] - if run == "sync": - result = getattr(fs._sync_fs, name)(*call_args, **kwargs) - else: - result = getattr(fs, name)(*call_args, **kwargs) - if asyncio.iscoroutine(result): - result = await await_bounded(result) - outcomes.append((comparable(result), sorted(map(repr, requests)), sorted(store))) + callback = _RecordingCallback() + call_kwargs = {k: callback if v == "recording" else v for k, v in kwargs.items()} + try: + if run == "sync": + result = getattr(fs._sync_fs, name)(*call_args, **call_kwargs) + else: + result = getattr(fs, name)(*call_args, **call_kwargs) + if asyncio.iscoroutine(result): + result = await await_bounded(result) + except Exception as e: + result = type(e) + downloaded = {p.name: p.read_bytes() for p in (tmp_path / run).iterdir()} + outcomes.append( + ( + comparable(result), + sorted(map(repr, requests)), + sorted(store), + downloaded, + callback.events, + ) + ) assert outcomes[0][1], "no request was sent" assert outcomes[1] == outcomes[0] From f8785188e095a37e98b4c23da566272e1727d05e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 10:31:49 +0900 Subject: [PATCH 17/19] Drop the test watchdog (GH-1239) pytest cancels a faulthandler timer when it reports a failure, before the teardown that the watchdog was meant to cover, and a timer left armed by one test could end a slow session fixture of the next. A thread that a regression blocks forever needs a process-level timeout, which the suite does not have. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 5a501be96..8b937f633 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1,7 +1,6 @@ import array import asyncio import contextlib -import faulthandler import gc import gzip import io @@ -51,23 +50,6 @@ wait_until, ) -# A regression that leaves a thread blocked can hang a test in its body or in -# its teardown, where no timeout of the test can end it; the watchdog dumps the -# stacks and ends the process instead, which fails the test. -WATCHDOG_TIMEOUT = 300 - - -@pytest.fixture(autouse=True) -def watchdog(): - # Rearmed for each test, so that it also covers the teardown. - faulthandler.dump_traceback_later(WATCHDOG_TIMEOUT, exit=True) - - -@pytest.fixture(autouse=True, scope="module") -def watchdog_cancellation(): - yield - faulthandler.cancel_dump_traceback_later() - class _RecordingCallback(Callback): """A progress callback that records its updates.""" From 888f55db4eaf355b7a27fffe1e197be222ca3be5 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 10:48:29 +0900 Subject: [PATCH 18/19] Fail the part copy once the other part runs in a waiting test (GH-1239) test_copy_object_with_multipart_upload_waits_for_running_parts let part 2 fail before part 1 had started, which then cancelled part 1 instead of waiting for it; it failed about 1 in 6 runs under load and once in the live run. Part 2 now fails once part 1 is running. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 8b937f633..bd150180b 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -197,16 +197,20 @@ async def test_copy_object_with_multipart_upload_waits_for_running_parts(self): ) ) events = [] + started = threading.Event() failed = threading.Event() def upload_part_copy(**kw): part_number = kw["part_number"] events.append(f"start {part_number}") if part_number == 1: - failed.wait(5) + started.set() + failed.wait(EVENT_TIMEOUT) time.sleep(0.05) events.append("end 1") return SimpleNamespace(etag='"e"', part_number=part_number) + # Fails once part 1 is running, which the copy cannot cancel. + assert started.wait(EVENT_TIMEOUT) failed.set() raise OSError("part failed") From da1f35fb566efc781f8498b64580707827aca8c0 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 10:59:10 +0900 Subject: [PATCH 19/19] Keep the running part of the abort test until the failure stops it (GH-1239) Part 1 of test_copy_object_with_multipart_upload_waits_for_running_parts ended a fixed 50 ms after part 2 failed, so a delayed failure handling let it free its worker and part 3 start on correct code. It now retries until the failure stops it, then ends a little later, and the test also checks that it was stopped. test_close_cancelled bounds its held completion. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 26 +++++++++++++--------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index bd150180b..cc603331a 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -197,21 +197,24 @@ async def test_copy_object_with_multipart_upload_waits_for_running_parts(self): ) ) events = [] - started = threading.Event() - failed = threading.Event() + # Part 1 retries until the failure of part 2 stops it, so that it + # runs until the copy handles the failure. + operation = ThrottledOperation(operation_name="UploadPartCopy") + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) def upload_part_copy(**kw): part_number = kw["part_number"] events.append(f"start {part_number}") if part_number == 1: - started.set() - failed.wait(EVENT_TIMEOUT) - time.sleep(0.05) - events.append("end 1") - return SimpleNamespace(etag='"e"', part_number=part_number) + try: + return throttled() + finally: + # Ends a little after the stop, so that an abort that + # does not wait for it is recorded first. + time.sleep(0.05) + events.append("end 1") # Fails once part 1 is running, which the copy cannot cancel. - assert started.wait(EVENT_TIMEOUT) - failed.set() + assert operation.started.wait(EVENT_TIMEOUT) raise OSError("part failed") sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) @@ -237,6 +240,7 @@ def upload_part_copy(**kw): # Part 3 waits for a worker and is not started after the failure. assert sorted(events[:2]) == ["start 1", "start 2"] assert events[2:] == ["end 1", "abort"] + assert operation.calls == 1 sync_fs.core.complete_multipart_upload.assert_not_called() @pytest.mark.asyncio @@ -2654,8 +2658,8 @@ async def test_close_cancelled(self, fs): def complete_multipart_upload(*args, **kwargs): started.set() - # The finally block of the test always releases it. - release.wait() + # The finally block of the test releases it. + release.wait(EVENT_TIMEOUT) completed.set() fs.core.complete_multipart_upload = mock.MagicMock(side_effect=complete_multipart_upload)