FabledCurator/tests/fake_s3.py

"""An in-process S3 bucket for round-trip tests: put, multipart upload,
get, and ListObjectsV2, kept in memory. It doesn't check signatures; the
signing has its own tests against AWS's published vectors. Like Fabled
Archives, it refuses to replace an object that is already there.

    with FakeS3() as s3:
        target = s3.target()
        ...
        s3.objects  # {key: bytes}
"""
from __future__ import annotations

import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, unquote, urlsplit
from xml.sax.saxutils import escape

from backend.app.services.backup_s3 import S3Target

BUCKET = "curator"


class FakeS3:
    def __init__(self):
        self.objects: dict[str, bytes] = {}
        self.uploads: dict[str, dict[int, bytes]] = {}
        outer = self

        class Handler(BaseHTTPRequestHandler):
            def _parts(self):
                u = urlsplit(self.path)
                path = unquote(u.path).lstrip("/")
                bucket, _, key = path.partition("/")
                query = parse_qs(u.query, keep_blank_values=True)   # `?uploads` has no value
                return bucket, key, {k: v[0] for k, v in query.items()}

            def _reply(self, status, body=b"", headers=None):
                self.send_response(status)
                for name, value in (headers or {}).items():
                    self.send_header(name, value)
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                if self.command != "HEAD":
                    self.wfile.write(body)

            def _body(self):
                length = int(self.headers.get("Content-Length") or 0)
                return self.rfile.read(length) if length else b""

            def do_PUT(self):
                _, key, q = self._parts()
                body = self._body()
                if "uploadId" in q:
                    outer.uploads[q["uploadId"]][int(q["partNumber"])] = body
                    return self._reply(200, headers={"ETag": f'"p{q["partNumber"]}"'})
                if key in outer.objects:
                    return self._reply(403, b"<Error><Code>AccessDenied</Code></Error>")
                outer.objects[key] = body
                return self._reply(200)

            def do_POST(self):
                _, key, q = self._parts()
                self._body()
                if "uploads" in q:
                    upload_id = f"u{len(outer.uploads) + 1}"
                    outer.uploads[upload_id] = {}
                    return self._reply(200, (
                        "<InitiateMultipartUploadResult>"
                        f"<UploadId>{upload_id}</UploadId>"
                        "</InitiateMultipartUploadResult>").encode())
                parts = outer.uploads.pop(q["uploadId"])
                outer.objects[key] = b"".join(parts[n] for n in sorted(parts))
                return self._reply(200, b"<CompleteMultipartUploadResult/>")

            def do_DELETE(self):
                _, _, q = self._parts()
                outer.uploads.pop(q.get("uploadId", ""), None)
                return self._reply(204)

            def do_GET(self):
                _, key, q = self._parts()
                if q.get("list-type") == "2":
                    prefix = q.get("prefix", "")
                    items = "".join(
                        f"<Contents><Key>{escape(k)}</Key><Size>{len(v)}</Size>"
                        "<LastModified>2026-01-01T00:00:00Z</LastModified></Contents>"
                        for k, v in sorted(outer.objects.items()) if k.startswith(prefix)
                    )
                    return self._reply(200, f"<ListBucketResult>{items}</ListBucketResult>".encode())
                if key not in outer.objects:
                    return self._reply(404, b"<Error><Code>NoSuchKey</Code></Error>")
                return self._reply(200, outer.objects[key])

            def log_message(self, *args):
                pass

        self.httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True)

    def __enter__(self):
        self.thread.start()
        return self

    def __exit__(self, *exc):
        self.httpd.shutdown()
        self.httpd.server_close()

    def target(self) -> S3Target:
        host, port = self.httpd.server_address
        return S3Target(f"http://{host}:{port}", BUCKET, "AK", "SK")