chore: import upstream snapshot with attribution

python/tests/integration/s3_mock.py

@@ -0,0 +1,272 @@
+# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
+
+import hashlib
+from collections import Counter
+from contextlib import ExitStack
+from datetime import datetime, timezone
+from unittest.mock import patch
+
+LAST_MODIFIED = datetime(2026, 3, 31, tzinfo=timezone.utc)
+
+_CORE_MODULES = [
+    "mirage.core.s3.read",
+    "mirage.core.s3.write",
+    "mirage.core.s3.stat",
+    "mirage.core.s3.readdir",
+    "mirage.core.s3.find",
+    "mirage.core.s3.du",
+    "mirage.core.s3.stream",
+    "mirage.core.s3.copy",
+    "mirage.core.s3.rename",
+    "mirage.core.s3.unlink",
+    "mirage.core.s3.rmdir",
+    "mirage.core.s3.rm",
+    "mirage.core.s3.mkdir",
+    "mirage.core.s3.create",
+    "mirage.core.s3.truncate",
+]
+
+
+class _AsyncMockBody:
+
+    def __init__(self, data: bytes) -> None:
+        self._data = data
+
+    async def read(self) -> bytes:
+        return self._data
+
+    async def iter_chunks(self, chunk_size: int = 8192):
+        for i in range(0, len(self._data), chunk_size):
+            yield self._data[i:i + chunk_size]
+
+
+def _mock_s3_error(code: str) -> Exception:
+    exc = Exception(code)
+    exc.response = {"Error": {"Code": code}}
+    return exc
+
+
+def _paginate_directory(objects, prefix):
+    common_prefixes: set[str] = set()
+    contents: list[dict[str, object]] = []
+    for key, data in sorted(objects.items()):
+        if not key.startswith(prefix):
+            continue
+        relative = key[len(prefix):]
+        if not relative:
+            contents.append({
+                "Key": key,
+                "Size": len(data),
+                "LastModified": LAST_MODIFIED
+            })
+            continue
+        if "/" in relative:
+            child = relative.split("/", 1)[0]
+            common_prefixes.add(prefix + child + "/")
+            continue
+        contents.append({
+            "Key": key,
+            "Size": len(data),
+            "LastModified": LAST_MODIFIED
+        })
+    return {
+        "CommonPrefixes": [{
+            "Prefix": v
+        } for v in sorted(common_prefixes)],
+        "Contents": contents,
+    }
+
+
+def _paginate_flat(objects, prefix):
+    return {
+        "Contents": [{
+            "Key": k,
+            "Size": len(v),
+            "LastModified": LAST_MODIFIED
+        } for k, v in sorted(objects.items()) if k.startswith(prefix)]
+    }
+
+
+def _slice_range(data: bytes, range_spec: str) -> bytes:
+    if not range_spec.startswith("bytes="):
+        return data
+    bounds = range_spec.removeprefix("bytes=").split("-", 1)
+    start = int(bounds[0]) if bounds[0] else 0
+    end = int(bounds[1]) if bounds[1] else len(data) - 1
+    return data[start:end + 1]
+
+
+class _MultiBucketPaginator:
+
+    def __init__(self, buckets: dict[str, dict[str, bytes]]) -> None:
+        self.buckets = buckets
+
+    async def paginate(self,
+                       Bucket: str,
+                       Prefix: str = "",
+                       Delimiter: str | None = None):
+        objects = self.buckets.get(Bucket, {})
+        if Delimiter == "/":
+            yield _paginate_directory(objects, Prefix)
+        else:
+            yield _paginate_flat(objects, Prefix)
+
+
+class MultiBucketS3Client:
+
+    def __init__(self,
+                 buckets: dict[str, dict[str, bytes]],
+                 versioned: set[str] | None = None,
+                 etag_suffix: str = "") -> None:
+        self.buckets = buckets
+        self.versioned = versioned or set()
+        self._versions: dict[tuple[str, str], list[tuple[str, bytes]]] = {}
+        # Non-empty simulates multipart-upload ETags ("<md5>-2"), which are
+        # NOT the MD5 of the content.
+        self.etag_suffix = etag_suffix
+        self.calls: Counter[str] = Counter()
+
+    def _etag(self, data: bytes) -> str:
+        return hashlib.md5(data).hexdigest() + self.etag_suffix
+
+    def _objects(self, bucket: str) -> dict[str, bytes]:
+        if bucket not in self.buckets:
+            self.buckets[bucket] = {}
+        return self.buckets[bucket]
+
+    def _track(self, bucket: str, key: str) -> str | None:
+        if bucket not in self.versioned:
+            return None
+        current = self.buckets.get(bucket, {}).get(key)
+        if current is None:
+            return None
+        history = self._versions.setdefault((bucket, key), [])
+        if not history or history[-1][1] != current:
+            vid = f"v{len(history) + 1}-{hashlib.md5(current).hexdigest()[:8]}"
+            history.append((vid, current))
+        return history[-1][0]
+
+    async def get_object(self,
+                         Bucket: str,
+                         Key: str,
+                         Range: str | None = None,
+                         VersionId: str | None = None) -> dict:
+        self.calls["get_object"] += 1
+        vid_for_resp = self._track(Bucket, Key)
+        if VersionId is not None:
+            history = self._versions.get((Bucket, Key), [])
+            for vid, data in history:
+                if vid == VersionId:
+                    vid_for_resp = vid
+                    break
+            else:
+                raise _mock_s3_error("NoSuchVersion")
+        else:
+            objects = self._objects(Bucket)
+            if Key not in objects:
+                raise _mock_s3_error("NoSuchKey")
+            data = objects[Key]
+        etag = self._etag(data)
+        if Range is not None:
+            data = _slice_range(data, Range)
+        resp: dict = {"Body": _AsyncMockBody(data), "ETag": f'"{etag}"'}
+        if vid_for_resp is not None:
+            resp["VersionId"] = vid_for_resp
+        return resp
+
+    async def head_object(self, Bucket: str, Key: str) -> dict:
+        self.calls["head_object"] += 1
+        objects = self._objects(Bucket)
+        if Key not in objects:
+            raise _mock_s3_error("NoSuchKey")
+        data = objects[Key]
+        etag = self._etag(data)
+        vid = self._track(Bucket, Key)
+        resp: dict = {
+            "ContentLength": len(data),
+            "LastModified": LAST_MODIFIED,
+            "ETag": f'"{etag}"',
+        }
+        if vid is not None:
+            resp["VersionId"] = vid
+        return resp
+
+    def get_paginator(self, name: str):
+        assert name == "list_objects_v2"
+        return _MultiBucketPaginator(self.buckets)
+
+    async def put_object(self, Bucket: str, Key: str, Body: bytes) -> None:
+        self._objects(Bucket)[Key] = Body
+
+    async def delete_object(self, Bucket: str, Key: str) -> None:
+        self._objects(Bucket).pop(Key, None)
+
+    async def copy_object(self, Bucket: str, CopySource: dict,
+                          Key: str) -> None:
+        src_bucket = CopySource.get("Bucket", Bucket)
+        src_key = CopySource["Key"]
+        src_objects = self._objects(src_bucket)
+        if src_key in src_objects:
+            self._objects(Bucket)[Key] = src_objects[src_key]
+
+    async def delete_objects(self, Bucket: str, Delete: dict) -> None:
+        objects = self._objects(Bucket)
+        for obj in Delete.get("Objects", []):
+            objects.pop(obj["Key"], None)
+
+    async def list_objects_v2(self,
+                              Bucket: str,
+                              Prefix: str = "",
+                              Delimiter: str = "",
+                              MaxKeys: int = 1000,
+                              **kwargs) -> dict:
+        del MaxKeys, kwargs
+        objects = self._objects(Bucket)
+        if Delimiter == "/":
+            return _paginate_directory(objects, Prefix)
+        return _paginate_flat(objects, Prefix)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        pass
+
+
+class MultiBucketSession:
+
+    def __init__(self,
+                 buckets: dict[str, dict[str, bytes]],
+                 versioned: set[str] | None = None,
+                 etag_suffix: str = "") -> None:
+        self._client = MultiBucketS3Client(buckets,
+                                           versioned=versioned,
+                                           etag_suffix=etag_suffix)
+
+    def client(self, **kwargs):
+        return self._client
+
+
+def patch_s3_session(session: MultiBucketSession) -> ExitStack:
+    stack = ExitStack()
+    for mod in _CORE_MODULES:
+        stack.enter_context(patch(f"{mod}.async_session",
+                                  return_value=session))
+    return stack
+
+
+def patch_s3_multi(buckets: dict[str, dict[str, bytes]],
+                   versioned: set[str] | None = None) -> ExitStack:
+    return patch_s3_session(MultiBucketSession(buckets, versioned=versioned))