notebook-snapshot-extension/snapshot/tests/test_storage_s3.py

"""Tests for the S3 storage adapter using a monkey-patched boto3."""

import pytest

boto3 = pytest.importorskip("boto3")
from botocore.exceptions import ClientError  # noqa: E402

from snapshot.storage import S3ConnectionError, S3Storage
from snapshot.storage.base import StorageAdapter


class _FakeS3Client:
    """In-memory S3 client sufficient for exercising S3Storage."""

    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self._objects = {}
        self._buckets = set()
        # Failure injection: if set, these are raised instead of the normal
        # behavior. Used to exercise the S3ConnectionError classification.
        self.head_bucket_error: Exception | None = None
        self.create_bucket_error: Exception | None = None

    def head_bucket(self, Bucket):
        if self.head_bucket_error is not None:
            raise self.head_bucket_error
        if Bucket not in self._buckets:
            raise ClientError({"Error": {"Code": "404"}}, "HeadBucket")

    def create_bucket(self, Bucket):
        if self.create_bucket_error is not None:
            raise self.create_bucket_error
        self._buckets.add(Bucket)

    def put_object(self, Bucket, Key, Body):
        self._buckets.add(Bucket)
        if isinstance(Body, bytes):
            data = Body
        else:
            data = Body.read()
        self._objects[(Bucket, Key)] = data

    def get_object(self, Bucket, Key):
        self.head_bucket(Bucket)
        if (Bucket, Key) not in self._objects:
            raise ClientError({"Error": {"Code": "NoSuchKey"}}, "GetObject")
        return {"Body": _Body(self._objects[(Bucket, Key)])}

    def delete_object(self, Bucket, Key):
        self._objects.pop((Bucket, Key), None)

    def head_object(self, Bucket, Key):
        self.head_bucket(Bucket)
        if (Bucket, Key) not in self._objects:
            raise ClientError({"Error": {"Code": "404"}}, "HeadObject")

    def list_objects_v2(self, Bucket, Prefix="", ContinuationToken=None):
        self.head_bucket(Bucket)
        keys = sorted(
            key for (b, key) in self._objects if b == Bucket and key.startswith(Prefix)
        )
        return {"Contents": [{"Key": key} for key in keys], "IsTruncated": False}


class _Body:
    def __init__(self, data: bytes):
        self._data = data

    def read(self):
        return self._data


def _make_adapter(monkeypatch, prefix=""):
    monkeypatch.setattr(boto3, "client", lambda *args, **kwargs: _FakeS3Client(**kwargs))
    return S3Storage(
        endpoint="localhost:9000",
        access_key="access",
        secret_key="secret",
        bucket="test-bucket",
        prefix=prefix,
    )


def _contract(adapter: StorageAdapter) -> None:
    adapter.put("a/b/file1.txt", b"hello")
    adapter.put("a/b/file2.txt", b"world")
    adapter.put("c/file3.txt", b"!")

    assert adapter.get("a/b/file1.txt") == b"hello"
    assert adapter.exists("a/b/file1.txt")
    assert not adapter.exists("a/b/missing.txt")

    with pytest.raises(KeyError):
        adapter.get("a/b/missing.txt")

    assert adapter.list_keys() == [
        "a/b/file1.txt",
        "a/b/file2.txt",
        "c/file3.txt",
    ]
    assert adapter.list_keys("a/b") == [
        "a/b/file1.txt",
        "a/b/file2.txt",
    ]

    adapter.delete("a/b/file1.txt")
    adapter.delete("a/b/file1.txt")  # idempotent
    assert not adapter.exists("a/b/file1.txt")
    assert adapter.list_keys() == [
        "a/b/file2.txt",
        "c/file3.txt",
    ]


def test_contract(monkeypatch):
    """Validate the adapter contract against the fake S3 client."""
    adapter = _make_adapter(monkeypatch)
    _contract(adapter)


def test_prefix_stripping(monkeypatch):
    """list_keys must strip the configured prefix from returned keys."""
    adapter = _make_adapter(monkeypatch, prefix="notebooks/")
    adapter.put("version1.ipynb", b"v1")
    adapter.put("version2.ipynb", b"v2")
    assert adapter.list_keys() == ["version1.ipynb", "version2.ipynb"]
    assert adapter.get("version1.ipynb") == b"v1"


def test_bucket_auto_creation(monkeypatch):
    """The bucket must be created automatically if it does not exist."""
    adapter = _make_adapter(monkeypatch)
    adapter.put("first.bin", b"x")
    assert adapter.exists("first.bin")


def test_imports_work_without_boto3():
    """Module-level imports must succeed even when boto3 is unavailable."""
    import snapshot
    import snapshot.storage
    import snapshot.storage.s3

    assert snapshot.storage.s3.S3Storage is not None
    assert snapshot.storage.StorageAdapter is not None


def test_lazy_import_error(monkeypatch):
    """Instantiation must raise ImportError when boto3 is absent."""
    real_import = __builtins__.__import__ if hasattr(__builtins__, "__import__") else __builtins__["__import__"]

    def fake_import(name, *args, **kwargs):
        if name == "boto3" or name.startswith("boto3."):
            raise ImportError(f"No module named {name!r}")
        return real_import(name, *args, **kwargs)

    monkeypatch.setattr("builtins.__import__", fake_import)
    with pytest.raises(ImportError):
        S3Storage(
            endpoint="localhost:9000",
            access_key="a",
            secret_key="s",
            bucket="b",
        )


# ---- connectivity / credentials checks at startup ------------------------


def _make_with_fake(monkeypatch, fake: _FakeS3Client) -> S3Storage:
    monkeypatch.setattr(boto3, "client", lambda *a, **kw: fake)
    return S3Storage(
        endpoint="localhost:9000",
        access_key="access",
        secret_key="secret",
        bucket="test-bucket",
    )


def test_permission_denied_raises_s3_connection_error(monkeypatch):
    """head_bucket returning 403/AccessDenied must surface as S3ConnectionError."""
    fake = _FakeS3Client()
    fake.head_bucket_error = ClientError(
        {"Error": {"Code": "AccessDenied"}}, "HeadBucket"
    )
    with pytest.raises(S3ConnectionError) as excinfo:
        _make_with_fake(monkeypatch, fake)
    msg = str(excinfo.value)
    assert "AccessDenied" in msg
    assert "localhost:9000" in msg
    assert "test-bucket" in msg
    assert "s3:ListBucket" in msg or "permission" in msg.lower()


def test_create_bucket_failure_raises_s3_connection_error(monkeypatch):
    """Bucket missing AND create_bucket failing must surface as S3ConnectionError."""
    fake = _FakeS3Client()
    fake.create_bucket_error = ClientError(
        {"Error": {"Code": "AccessDenied"}}, "CreateBucket"
    )
    with pytest.raises(S3ConnectionError) as excinfo:
        _make_with_fake(monkeypatch, fake)
    msg = str(excinfo.value)
    assert "does not exist" in msg
    assert "could not be created" in msg
    assert "CreateBucket" in msg


def test_endpoint_unreachable_raises_s3_connection_error(monkeypatch):
    """Network failures (e.g. EndpointConnectionError) surface as S3ConnectionError."""
    from botocore.exceptions import EndpointConnectionError

    fake = _FakeS3Client()
    fake.head_bucket_error = EndpointConnectionError(
        endpoint_url="http://localhost:9000"
    )
    with pytest.raises(S3ConnectionError) as excinfo:
        _make_with_fake(monkeypatch, fake)
    msg = str(excinfo.value)
    assert "Cannot connect" in msg
    assert "localhost:9000" in msg
    assert "EndpointConnectionError" in msg


def test_credentials_error_surfaces_credentials_hint(monkeypatch):
    """Credential-related boto3 errors mention the S3_* env vars in the hint."""
    from botocore.exceptions import NoCredentialsError

    fake = _FakeS3Client()
    fake.head_bucket_error = NoCredentialsError()
    with pytest.raises(S3ConnectionError) as excinfo:
        _make_with_fake(monkeypatch, fake)
    msg = str(excinfo.value)
    assert "S3_ACCESS_KEY" in msg
    assert "S3_SECRET_KEY" in msg


def test_connectivity_succeeds_silently_when_bucket_exists(monkeypatch):
    """If head_bucket succeeds (bucket exists), no create is called."""
    fake = _FakeS3Client()
    fake._buckets.add("test-bucket")  # simulate pre-existing bucket
    create_called = {"flag": False}
    original_create = fake.create_bucket

    def tracking_create(Bucket):
        create_called["flag"] = True
        return original_create(Bucket)

    fake.create_bucket = tracking_create
    _make_with_fake(monkeypatch, fake)
    assert create_called["flag"] is False, "create_bucket must not be called when bucket already exists"