feat userver: add mutlipart methods support for testsuite s3api plugin

basil-shuman · basil-shuman · commit 9c711cd745fa · 2025-12-02T14:38:02.000+03:00
commit_hash:8846c5dddd8f795c73d2c57e611c3dbc61d5148f
diff --git a/.mapping.json b/.mapping.json
@@ -4814,6 +4814,9 @@
   "testsuite/tests/metrics/CMakeLists.txt":"taxi/uservices/userver/testsuite/tests/metrics/CMakeLists.txt",
   "testsuite/tests/metrics/conftest.py":"taxi/uservices/userver/testsuite/tests/metrics/conftest.py",
   "testsuite/tests/metrics/test_metrics.py":"taxi/uservices/userver/testsuite/tests/metrics/test_metrics.py",
+  "testsuite/tests/s3api/CMakeLists.txt":"taxi/uservices/userver/testsuite/tests/s3api/CMakeLists.txt",
+  "testsuite/tests/s3api/conftest.py":"taxi/uservices/userver/testsuite/tests/s3api/conftest.py",
+  "testsuite/tests/s3api/test_multipart_upload.py":"taxi/uservices/userver/testsuite/tests/s3api/test_multipart_upload.py",
   "testsuite/tests/sql_coverage_empty/CMakeLists.txt":"taxi/uservices/userver/testsuite/tests/sql_coverage_empty/CMakeLists.txt",
   "testsuite/tests/sql_coverage_empty/conftest.py":"taxi/uservices/userver/testsuite/tests/sql_coverage_empty/conftest.py",
   "testsuite/tests/sql_coverage_empty/test_coverage_empty.py":"taxi/uservices/userver/testsuite/tests/sql_coverage_empty/test_coverage_empty.py",
diff --git a/testsuite/pytest_plugins/pytest_userver/plugins/s3api.py b/testsuite/pytest_plugins/pytest_userver/plugins/s3api.py
@@ -34,11 +34,23 @@ def _mock_all(request):
             return mock_impl.get_object(request)
 
         if request.method == 'PUT':
+            if 'uploadId' in request.query:
+                if 'partNumber' in request.query:
+                    return mock_impl.upload_part(request)
+                return mockserver.make_response('Unknown or unsupported method', 404)
             if request.headers.get('x-amz-copy-source', None):
                 return mock_impl.copy_object(request)
             return mock_impl.put_object(request)
 
+        if request.method == 'POST':
+            if 'uploadId' in request.query:
+                return mock_impl.complete_multipart_upload(request)
+            if 'uploads' in request.query:
+                return mock_impl.create_multipart_upload(request)
+
         if request.method == 'DELETE':
+            if 'uploadId' in request.query:
+                return mock_impl.abort_multipart_upload(request)
             return mock_impl.delete_object(request)
 
         if request.method == 'HEAD':
diff --git a/testsuite/pytest_plugins/pytest_userver/s3api.py b/testsuite/pytest_plugins/pytest_userver/s3api.py
@@ -1,13 +1,167 @@
+import collections
 from collections.abc import Mapping
+from collections.abc import MutableMapping
 import dataclasses
 import datetime as dt
 import hashlib
+import os
 import pathlib
 import sys
+import xml.etree.ElementTree
 
 import dateutil.tz as tz
 
 
+class _S3NoSuchUploadError(Exception):
+    code = 'NoSuchUpload'
+    message = 'The specified multipart upload does not exist.'
+
+    def __str__(self):
+        return _S3NoSuchUploadError.message
+
+
+class _S3InvalidPartError(Exception):
+    code = 'InvalidPart'
+    message = (
+        'One or more of the specified parts could not be found.'
+        ' The part might not have been uploaded, or the specified'
+        " ETag might not have matched the uploaded part's ETag."
+    )
+
+    def __str__(self):
+        return _S3InvalidPartError.message
+
+
+class _S3InvalidPartOrderError(Exception):
+    code = 'InvalidPartOrder'
+    message = 'The list of parts was not in ascending order. The parts list must be specified in order by part number.'
+
+    def __str__(self):
+        return _S3InvalidPartOrderError.message
+
+
+class _S3EntityTooSmallError(Exception):
+    code = 'EntityTooSmall'
+    message = 'Your proposed upload is smaller than the minimum allowed object size.'
+
+    def __str__(self):
+        return _S3EntityTooSmallError.message
+
+
+class _S3ClientError(Exception):
+    def __init__(self, msg: str):
+        self._msg = msg
+
+    def __str__(self):
+        return self._msg
+
+
+@dataclasses.dataclass
+class _S3UploadPart:
+    data: bytearray
+    meta: Mapping[str, str]
+
+
+@dataclasses.dataclass
+class _S3Upload:
+    parts: MutableMapping[int, _S3UploadPart]
+    meta: Mapping[str, str]
+
+
+@dataclasses.dataclass
+class _S3BucketUploadStorage:
+    def __init__(self):
+        self._storage: dict[str, _S3Upload] = {}
+
+    @staticmethod
+    def _generate_upload_id():
+        return os.urandom(15).hex()
+
+    @staticmethod
+    def _generate_etag(data):
+        return hashlib.md5(data).hexdigest()
+
+    def create_multipart_upload(self, key: str, user_defined_meta: Mapping[str, str] | None = None):
+        key_path = pathlib.Path(key)
+        upload_id = _S3BucketUploadStorage._generate_upload_id()
+
+        upload_meta = {
+            'Key': str(key_path),
+            'UploadId': upload_id,
+        }
+
+        if user_defined_meta:
+            upload_meta.update(user_defined_meta)
+
+        self._storage[upload_id] = _S3Upload(parts={}, meta=upload_meta)
+        return upload_meta
+
+    def abort_multipart_uplod(self, key: str, upload_id: str):
+        key_path = pathlib.Path(key)
+        upload = self._storage.get(upload_id)
+        if not upload or upload.meta['Key'] != str(key_path):
+            raise _S3NoSuchUploadError()
+        return self._storage.pop(upload_id)
+
+    def upload_part(
+        self,
+        key: str,
+        upload_id: str,
+        part_number: int,
+        data: bytearray,
+        last_modified: dt.datetime | str | None = None,
+    ):
+        if part_number < 1 or part_number > 10000:
+            raise _S3ClientError('partNumber value is expected to be between 1 and 10000')
+
+        key_path = pathlib.Path(key)
+        upload = self._storage.get(upload_id)
+        if not upload or upload.meta['Key'] != str(key_path):
+            raise _S3NoSuchUploadError()
+
+        if last_modified is None:
+            # Timezone is needed for RFC 3339 timeformat used by S3
+            last_modified = dt.datetime.now().replace(tzinfo=tz.tzlocal()).isoformat()
+        elif isinstance(last_modified, dt.datetime):
+            last_modified = last_modified.isoformat()
+
+        meta = {
+            'ETag': self._generate_etag(data),
+            'Last-Modified': last_modified,
+            'Size': str(sys.getsizeof(data)),
+        }
+
+        new_part = _S3UploadPart(data, meta)
+        upload.parts[part_number] = new_part
+        return new_part
+
+    def complete_multipart_upload(self, key: str, upload_id: str, parts_to_complete: list):
+        key_path = pathlib.Path(key)
+        upload = self._storage.get(upload_id)
+
+        if not upload or upload.meta['Key'] != str(key_path):
+            raise _S3NoSuchUploadError()
+
+        uploaded_parts = sorted(
+            ({'PartNumber': part_number, 'ETag': info.meta['ETag']} for part_number, info in upload.parts.items()),
+            key=lambda item: item['PartNumber'],
+        )
+        if uploaded_parts != parts_to_complete:
+            raise _S3InvalidPartOrderError()
+
+        merged_data = bytearray()
+        for part in parts_to_complete:
+            part_number = part['PartNumber']
+            uploded_part = upload.parts[part_number]
+            merged_data += uploded_part.data
+
+        if not merged_data:
+            raise _S3EntityTooSmallError()
+
+        self._storage.pop(upload_id)
+        return {'Data': merged_data, 'Upload': upload}
+
+
 @dataclasses.dataclass
 class S3Object:
     data: bytearray
@@ -70,10 +224,13 @@ def delete_object(self, key) -> S3Object | None:
 
 
 class S3HandleMock:
+    _s3_xml_nss = {'s3': 'http://s3.amazonaws.com/doc/2006-03-01/'}
+
     def __init__(self, mockserver, s3_mock_storage, mock_base_url):
         self._mockserver = mockserver
         self._base_url = mock_base_url
         self._storage = s3_mock_storage
+        self._uploads = collections.defaultdict(_S3BucketUploadStorage)
 
     def _get_bucket_name(self, request):
         return request.headers['Host'].split('.')[0]
@@ -137,6 +294,39 @@ def _generate_get_objects_xml(
             </ListBucketResult>
             """
 
+    @staticmethod
+    def _generate_error_response_xml(code: str, message: str, resource: str):
+        return (
+            '<?xml version="1.0" encoding="UTF-8"?>'
+            '<Error>'
+            f'<Code>{code}</Code>'
+            f'<Message>{message}</Message>'
+            f'<Resource>{resource}</Resource>'
+            f'<RequestId>{os.urandom(15).hex()}</RequestId>'
+            '</Error>'
+        )
+
+    @staticmethod
+    def _parse_complete_multipart_xml_body(request_body: str):
+        xml_root_node = xml.etree.ElementTree.fromstring(request_body)
+        if xml_root_node is None or xml_root_node.tag != f'{{{S3HandleMock._s3_xml_nss["s3"]}}}CompleteMultipartUpload':
+            raise _S3ClientError('missing CompleteMultipartUpload in request body')
+
+        parts_to_complete = []
+        for xml_part in xml_root_node.findall('s3:Part', S3HandleMock._s3_xml_nss):
+            xml_part_number = xml_part.find('s3:PartNumber', S3HandleMock._s3_xml_nss)
+            if xml_part_number is None or not xml_part_number.text:
+                raise _S3ClientError('missing CompleteMultipartUpload.Part.PartNumber')
+            part_number_value = int(xml_part_number.text)
+
+            xml_etag = xml_part.find('s3:ETag', S3HandleMock._s3_xml_nss)
+            if xml_etag is None or not xml_etag.text:
+                raise _S3ClientError('missing CompleteMultipartUpload.Part.ETag')
+
+            parts_to_complete.append({'ETag': xml_etag.text, 'PartNumber': part_number_value})
+
+        return parts_to_complete
+
     def get_object(self, request):
         key = self._extract_key(request)
 
@@ -161,7 +351,7 @@ def put_object(self, request):
         user_defined_meta = {}
         for meta_key, meta_value in request.headers.items():
             # https://docs.amazonaws.cn/en_us/AmazonS3/latest/userguide/UsingMetadata.html
-            if meta_key.startswith('x-amz-meta-'):
+            if meta_key.startswith('x-amz-meta-') or meta_key in ['Content-Type', 'Content-Disposition']:
                 user_defined_meta[meta_key] = meta_value
 
         meta = bucket_storage.put_object(key, data, user_defined_meta)
@@ -234,3 +424,109 @@ def get_object_head(self, request):
             200,
             headers=s3_object.meta,
         )
+
+    def create_multipart_upload(self, request):
+        key = self._extract_key(request)
+        bucket_name = self._get_bucket_name(request)
+        bucket_uploads = self._uploads[bucket_name]
+
+        user_defined_meta = {}
+        for meta_key, meta_value in request.headers.items():
+            # https://docs.amazonaws.cn/en_us/AmazonS3/latest/userguide/UsingMetadata.html
+            if meta_key.startswith('x-amz-meta-') or meta_key in ['Content-Type', 'Content-Disposition']:
+                user_defined_meta[meta_key] = meta_value
+
+        meta = bucket_uploads.create_multipart_upload(key, user_defined_meta)
+        response_body = (
+            '<?xml version="1.0" encoding="UTF-8"?>'
+            '<InitiateMultipartUploadResult>'
+            f'<Bucket>{bucket_name}</Bucket>'
+            f'<Key>{key}</Key>'
+            f'<UploadId>{meta["UploadId"]}</UploadId>'
+            '</InitiateMultipartUploadResult>'
+        )
+        # Some clients like AWS SDK for C++ parse not empty body as XML
+        return self._mockserver.make_response(response_body, 200)
+
+    def abort_multipart_upload(self, request):
+        key = self._extract_key(request)
+        upload_id = request.query['uploadId']
+        bucket_uploads = self._uploads[self._get_bucket_name(request)]
+        try:
+            bucket_uploads.abort_multipart_uplod(key, upload_id)
+        except _S3NoSuchUploadError as exc:
+            # https://docs.aws.amazon.com/AmazonS3/latest/API/API_AbortMultipartUpload.html
+            # #API_AbortMultipartUpload_Errors
+            response_body = S3HandleMock._generate_error_response_xml(
+                exc.code, exc.message, f'{request.path}?uploadId={upload_id}'
+            )
+            return self._mockserver.make_response(response_body, 404)
+
+        # Some clients like AWS SDK for C++ parse not empty body as XML
+        return self._mockserver.make_response('', 204)
+
+    def upload_part(self, request):
+        key = self._extract_key(request)
+        bucket_name = self._get_bucket_name(request)
+        upload_id = request.query['uploadId']
+        part_number = int(request.query['partNumber'])
+        bucket_uploads = self._uploads[bucket_name]
+        data = request.get_data()
+        try:
+            upload_part = bucket_uploads.upload_part(key, upload_id, part_number, data)
+        except _S3ClientError as exc:
+            return self._mockserver.make_response(str(exc), 400)
+        except _S3NoSuchUploadError as exc:
+            # https://docs.aws.amazon.com/AmazonS3/latest/API/API_UploadPart.html
+            response_body = S3HandleMock._generate_error_response_xml(
+                exc.code,
+                exc.message,
+                f'{request.path}?uploadId={upload_id}',
+            )
+            return self._mockserver.make_response(response_body, 404)
+
+        return self._mockserver.make_response(status=200, headers={'ETag': upload_part.meta['ETag']})
+
+    def complete_multipart_upload(self, request):
+        key = self._extract_key(request)
+        bucket_name = self._get_bucket_name(request)
+        bucket_uploads = self._uploads[bucket_name]
+        bucket_storage = self._storage[bucket_name]
+        upload_id = request.query['uploadId']
+        try:
+            parts_to_complete = S3HandleMock._parse_complete_multipart_xml_body(request.get_data().decode())
+            completed_upload = bucket_uploads.complete_multipart_upload(key, upload_id, parts_to_complete)
+        except _S3NoSuchUploadError as exc:
+            # https://docs.aws.amazon.com/AmazonS3/latest/API/API_CompleteMultipartUpload.html
+            response_body = S3HandleMock._generate_error_response_xml(
+                exc.code,
+                exc.message,
+                f'{request.path}?uploadId={upload_id}',
+            )
+            return self._mockserver.make_response(response_body, 404)
+        except (_S3InvalidPartError, _S3InvalidPartOrderError, _S3EntityTooSmallError) as exc:
+            # https://docs.aws.amazon.com/AmazonS3/latest/API/API_CompleteMultipartUpload.html
+            response_body = S3HandleMock._generate_error_response_xml(
+                exc.code,
+                exc.message,
+                f'{request.path}?uploadId={upload_id}',
+            )
+            return self._mockserver.make_response(response_body, 400)
+        except _S3ClientError as exc:
+            return self._mockserver.make_response(str(exc), 400)
+
+        meta = bucket_storage.put_object(key, completed_upload['Data'], completed_upload['Upload'].meta)
+        response_body = (
+            '<?xml version="1.0" encoding="UTF-8"?>'
+            '<CompleteMultipartUploadResult>'
+            f'<Location>{request.path}</Location>'
+            f'<Bucket>{bucket_name}</Bucket>'
+            f'<Key>{key}</Key>'
+            f'<ETag>{meta["ETag"]}</ETag>'
+            '</CompleteMultipartUploadResult>'
+        )
+        return self._mockserver.make_response(
+            response_body,
+            status=200,
+            headers=meta,
+        )
diff --git a/testsuite/tests/CMakeLists.txt b/testsuite/tests/CMakeLists.txt
@@ -3,6 +3,7 @@ project(userver-testsuite-tests)
 add_subdirectory(chaos)
 add_subdirectory(config)
 add_subdirectory(metrics)
+# add_subdirectory(s3api) TODO: TAXICOMMON-10381
 add_subdirectory(sql_coverage_empty)
 add_subdirectory(tskv)
 # add_subdirectory(sql_coverage_full) TODO: TAXICOMMON-10381
diff --git a/testsuite/tests/s3api/CMakeLists.txt b/testsuite/tests/s3api/CMakeLists.txt
@@ -0,0 +1,7 @@
+add_test(
+    NAME "${PROJECT_NAME}-s3api"
+    COMMAND "${TESTSUITE_PYTHON_BINARY}" -m pytest ${CMAKE_CURRENT_SOURCE_DIR}
+    WORKING_DIRECTORY "${USERVER_ROOT_DIR}/testsuite/pytest_plugins"
+)
+
+set_tests_properties("${PROJECT_NAME}-s3api" PROPERTIES TIMEOUT 600)
diff --git a/testsuite/tests/s3api/conftest.py b/testsuite/tests/s3api/conftest.py
@@ -0,0 +1 @@
+pytest_plugins = ['pytest_userver.plugins.core', 'pytest_userver.plugins.s3api']
diff --git a/testsuite/tests/s3api/test_multipart_upload.py b/testsuite/tests/s3api/test_multipart_upload.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+pytest_plugins = ['pytest_userver.plugins.core', 'pytest_userver.plugins.s3api']`