strukto-ai--mirage/python/tests/integration/s3_mock.py

# ========= 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))