fix: artifact_bucket

schedule/tests/test_worker.py

@@ -0,0 +1,154 @@
+"""Tests for NodeExecutor bucket-aware artifact download (P0-3).
+
+Pre-fix the worker's ``object_store`` was bound to the global version
+bucket at startup, so a workspace whose ``Workspaces.artifact_bucket``
+points at a custom S3 bucket would always 404 on download. The fix
+introduces a per-bucket store cache so the worker reads from whatever
+bucket the artifact actually lives in.
+
+These tests cover only the routing logic — no live S3 / MySQL.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+
+def _make_executor() -> tuple[SimpleNamespace, MagicMock, MagicMock]:
+    """Construct a NodeExecutor with the bucket-store cache pre-seeded.
+
+    Returns ``(executor, default_store, storage_client)``. The default
+    store is whatever the worker would receive via ``object_store=`` at
+    construction time; it must be returned untouched by ``_store_for`` for
+    the global version bucket.
+    """
+    from schedule.worker import NodeExecutor
+
+    default_store = MagicMock(name="default_store")
+    storage_client = MagicMock(name="storage_client")
+    executor = NodeExecutor(
+        session_factory=MagicMock(),
+        object_store=default_store,
+        storage_client=storage_client,
+    )
+    return executor, default_store, storage_client
+
+
+def test_store_for_default_bucket_returns_injected_store() -> None:
+    """The default version bucket must reuse the injected ``object_store``.
+
+    Otherwise the common (no-override) path would pay for an extra
+    ``create_storage`` call on every run — wasted work.
+    """
+    from common.config import settings
+
+    executor, default_store, _ = _make_executor()
+    store = executor._store_for(settings.s3_version_bucket)
+    assert store is default_store
+
+
+def test_store_for_custom_bucket_uses_build_factory(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A custom bucket_name must produce a store via ``build_object_store``."""
+    custom_store = MagicMock(name="custom_store")
+    with patch(
+        "schedule.service.build_object_store",
+        return_value=custom_store,
+    ) as mock_build:
+        executor, _, _ = _make_executor()
+        store = executor._store_for("my-workspace-artifacts")
+
+    assert store is custom_store
+    mock_build.assert_called_once_with(bucket_name="my-workspace-artifacts")
+
+
+def test_store_for_custom_bucket_is_cached(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Repeated lookups for the same custom bucket must hit the cache."""
+    custom_store = MagicMock(name="custom_store")
+    with patch(
+        "schedule.service.build_object_store",
+        return_value=custom_store,
+    ) as mock_build:
+        executor, _, _ = _make_executor()
+        executor._store_for("my-workspace-artifacts")
+        executor._store_for("my-workspace-artifacts")
+        executor._store_for("my-workspace-artifacts")
+
+    mock_build.assert_called_once_with(bucket_name="my-workspace-artifacts")
+
+
+@pytest.mark.asyncio
+async def test_download_artifact_uses_per_bucket_store() -> None:
+    """_download_artifact must call ``store.get(object_key)`` on the
+    bucket-bound store, NOT the default version-bucket store.
+
+    Pre-fix this would route every workspace's reads through the global
+    version bucket, silently 404'ing for any workspace that overrode
+    ``artifact_bucket``.
+    """
+    custom_store = MagicMock(name="custom_store")
+    custom_store.get = AsyncMock(return_value=b"hello world")
+    executor, default_store, _ = _make_executor()
+    # Pretend the cache already has the custom bucket wired up.
+    executor._bucket_stores["my-workspace-artifacts"] = custom_store
+
+    payload = b"hello world"
+    expected_hash = hashlib.sha256(payload).hexdigest()
+    result = await executor._download_artifact(
+        bucket_name="my-workspace-artifacts",
+        object_key="ws-1/user-1/script.py",
+        content_hash=expected_hash,
+    )
+
+    assert result == payload
+    custom_store.get.assert_awaited_once_with("ws-1/user-1/script.py")
+    # The default store must NOT have been touched — that was the bug.
+    default_store.get.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_download_artifact_uses_default_store_when_bucket_matches() -> None:
+    """Sanity: default bucket reads still flow through the injected store."""
+    from common.config import settings
+
+    default_store = MagicMock(name="default_store")
+    default_store.get = AsyncMock(return_value=b"payload")
+    executor = type(_make_executor()[0])(
+        session_factory=MagicMock(),
+        object_store=default_store,
+        storage_client=MagicMock(),
+    )
+    payload = b"payload"
+    result = await executor._download_artifact(
+        bucket_name=settings.s3_version_bucket,
+        object_key="ws-1/user-1/script.py",
+        content_hash=hashlib.sha256(payload).hexdigest(),
+    )
+
+    assert result == payload
+    default_store.get.assert_awaited_once_with("ws-1/user-1/script.py")
+
+
+@pytest.mark.asyncio
+async def test_download_artifact_rejects_hash_mismatch() -> None:
+    """Hash mismatch must raise regardless of which store was used —
+    protects against corrupted bytes landing in any bucket.
+    """
+    custom_store = MagicMock(name="custom_store")
+    custom_store.get = AsyncMock(return_value=b"corrupted")
+    executor, default_store, _ = _make_executor()
+    executor._bucket_stores["my-workspace-artifacts"] = custom_store
+
+    with pytest.raises(ValueError, match="hash mismatch"):
+        await executor._download_artifact(
+            bucket_name="my-workspace-artifacts",
+            object_key="ws-1/user-1/script.py",
+            content_hash="0" * 64,  # cannot match corrupted bytes
+        )
+
+    # Reads still went through the bucket-bound store, not the default.
+    custom_store.get.assert_awaited_once()
+    default_store.get.assert_not_called()