chore: import upstream snapshot with attribution

python/tests/workspace/test_snapshot_drift.py

@@ -0,0 +1,211 @@
+# ========= 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 asyncio
+from contextlib import ExitStack
+
+import pytest
+
+from mirage.resource.ram import RAMResource
+from mirage.resource.s3 import S3Config, S3Resource
+from mirage.types import DriftPolicy, MountMode
+from mirage.workspace import Workspace
+from mirage.workspace.snapshot import ContentDriftError, install_fingerprints
+from tests.integration.s3_mock import patch_s3_multi
+
+
+def _load(*args, **kwargs):
+    return asyncio.run(Workspace.load(*args, **kwargs))
+
+
+def test_install_fingerprints_pins_revision_and_queues_drift():
+    ws = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
+                   mode=MountMode.WRITE)
+    entries = [
+        {
+            "path": "/m/pinned.txt",
+            "mount_prefix": "/m/",
+            "revision": "v9"
+        },
+        {
+            "path": "/m/checked.txt",
+            "mount_prefix": "/m/",
+            "fingerprint": "fp1"
+        },
+    ]
+
+    install_fingerprints(ws, entries, DriftPolicy.STRICT)
+
+    mount = ws._registry.mount_for("/m/pinned.txt")
+    assert mount.revisions["/m/pinned.txt"] == "v9"
+    assert ws._drift_check_pending is True
+    queued = [path for (_, path, _) in ws._pending_drift]
+    assert "/m/checked.txt" in queued
+
+
+def _config() -> S3Config:
+    return S3Config(
+        bucket="test-bucket",
+        region="us-east-1",
+        aws_access_key_id="fake",
+        aws_secret_access_key="fake",
+    )
+
+
+def test_strict_load_raises_when_s3_etag_drifts(tmp_path):
+    """Snapshot a workspace with one read on /s3, mutate the underlying
+    object so its ETag changes, load with strict policy: the next read
+    must raise ContentDriftError rather than silently serve drifted
+    bytes.
+    """
+    store = {"data.csv": b"version 1 bytes\n"}
+    with ExitStack() as stack:
+        stack.enter_context(patch_s3_multi({"test-bucket": store}))
+        src = Workspace({"/s3": (S3Resource(_config()), MountMode.WRITE)},
+                        mode=MountMode.WRITE)
+        result = asyncio.run(src.execute("cat /s3/data.csv"))
+        assert b"version 1" in result.stdout
+
+        snap = tmp_path / "snap.tar"
+        asyncio.run(src.snapshot(snap))
+        store["data.csv"] = b"VERSION 2 DRIFTED\n"
+
+        dst = _load(snap, resources={"/s3": S3Resource(_config())})
+        with pytest.raises(ContentDriftError) as exc_info:
+            asyncio.run(dst.execute("cat /s3/data.csv"))
+        assert exc_info.value.path == "/s3/data.csv"
+        live = exc_info.value.live_fingerprint
+        recorded = exc_info.value.snapshot_fingerprint
+        assert live != recorded
+
+
+def test_off_load_serves_drifted_bytes_silently(tmp_path):
+    """drift_policy=OFF disables the check: the load returns the new
+    bytes with no error. This is the only opt-out from drift detection.
+    """
+    store = {"data.csv": b"version 1 bytes\n"}
+    with ExitStack() as stack:
+        stack.enter_context(patch_s3_multi({"test-bucket": store}))
+        src = Workspace({"/s3": (S3Resource(_config()), MountMode.WRITE)},
+                        mode=MountMode.WRITE)
+        asyncio.run(src.execute("cat /s3/data.csv"))
+
+        snap = tmp_path / "snap.tar"
+        asyncio.run(src.snapshot(snap))
+        store["data.csv"] = b"VERSION 2 DRIFTED\n"
+
+        dst = _load(snap,
+                    resources={"/s3": S3Resource(_config())},
+                    drift_policy=DriftPolicy.OFF)
+        result = asyncio.run(dst.execute("cat /s3/data.csv"))
+        assert b"VERSION 2 DRIFTED" in result.stdout
+
+
+def test_strict_load_passes_when_etag_unchanged(tmp_path):
+    """The control case: same bytes still in S3, strict load must succeed
+    and serve them. Verifies the check is precise (no false positives).
+    """
+    store = {"data.csv": b"stable bytes\n"}
+    with ExitStack() as stack:
+        stack.enter_context(patch_s3_multi({"test-bucket": store}))
+        src = Workspace({"/s3": (S3Resource(_config()), MountMode.WRITE)},
+                        mode=MountMode.WRITE)
+        asyncio.run(src.execute("cat /s3/data.csv"))
+
+        snap = tmp_path / "snap.tar"
+        asyncio.run(src.snapshot(snap))
+
+        dst = _load(snap, resources={"/s3": S3Resource(_config())})
+        result = asyncio.run(dst.execute("cat /s3/data.csv"))
+        assert b"stable bytes" in result.stdout
+
+
+def test_unrecorded_path_skips_drift_check(tmp_path):
+    """A path the agent did not read at snapshot time has no recorded
+    fingerprint, so it must NOT be drift-checked at load — just live-
+    served. Tests that drift-check is opt-in per recorded path.
+    """
+    store = {
+        "read-me.txt": b"recorded\n",
+        "added-later.txt": b"not in snapshot\n"
+    }
+    with ExitStack() as stack:
+        stack.enter_context(patch_s3_multi({"test-bucket": store}))
+        src = Workspace({"/s3": (S3Resource(_config()), MountMode.WRITE)},
+                        mode=MountMode.WRITE)
+        asyncio.run(src.execute("cat /s3/read-me.txt"))
+
+        snap = tmp_path / "snap.tar"
+        asyncio.run(src.snapshot(snap))
+
+        dst = _load(snap, resources={"/s3": S3Resource(_config())})
+        result = asyncio.run(dst.execute("cat /s3/added-later.txt"))
+        assert b"not in snapshot" in result.stdout
+
+
+def test_version_pin_serves_original_bytes_on_versioned_bucket(tmp_path):
+    """On a versioned S3 bucket, snapshot captures the object's VersionId
+    alongside its ETag. After the live bytes drift, load with default
+    STRICT policy pins reads to the recorded VersionId and serves the
+    ORIGINAL bytes the agent saw, instead of raising drift or returning
+    the current head.
+    """
+    store = {"data.csv": b"original\n"}
+    with ExitStack() as stack:
+        stack.enter_context(
+            patch_s3_multi({"test-bucket": store}, versioned={"test-bucket"}))
+        src = Workspace({"/s3": (S3Resource(_config()), MountMode.WRITE)},
+                        mode=MountMode.WRITE)
+        asyncio.run(src.execute("cat /s3/data.csv"))
+
+        snap = tmp_path / "snap.tar"
+        asyncio.run(src.snapshot(snap))
+        store["data.csv"] = b"mutated bytes\n"
+
+        dst = _load(snap, resources={"/s3": S3Resource(_config())})
+        # Cache holds snapshot bytes; clear so we hit S3 and exercise the
+        # pin path, not the cache path.
+        _drop_path_from_cache(dst, "/s3/data.csv")
+        result = asyncio.run(dst.execute("cat /s3/data.csv"))
+        assert result.stdout == b"original\n"
+
+
+def _drop_path_from_cache(ws, path: str) -> None:
+    cache = ws._cache
+    cache._entries.pop(path, None)
+    cache._store.files.pop(path, None)
+
+
+def test_live_only_mount_does_not_block_snapshot(tmp_path, caplog):
+    """Workspaces with non-SUPPORTS_SNAPSHOT mounts (RAM here as a stand-
+    in for Gmail/Slack/Linear) snapshot fine; no fingerprints are
+    captured for those paths and the load layer logs an honest warning
+    surfacing the live-only mount list.
+    """
+    src = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
+                    mode=MountMode.WRITE)
+    asyncio.run(src.execute("echo body > /m/note.txt"))
+    asyncio.run(src.execute("cat /m/note.txt"))
+
+    snap = tmp_path / "snap.tar"
+    asyncio.run(src.snapshot(snap))
+
+    with caplog.at_level("WARNING"):
+        _load(snap)
+    assert any("live-only" in r.message.lower()
+               or "live-only" in r.getMessage().lower()
+               for r in caplog.records) or any(
+                   "no drift" in r.message.lower()
+                   or "no drift" in r.getMessage().lower()
+                   for r in caplog.records)