model-platform/schedule/tests/test_worker.py

"""Tests for NodeExecutor bucket-aware artifact download (P0-3) and
user-status re-verification (P0-5 / C1).

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.

P0-5 / C1: the worker previously trusted the ``triggered_by`` user_id
without re-checking ``Users.status`` / ``is_deleted``. A user disabled
after a run was created would still have their schedules execute.
The fix adds ``_assert_user_active`` to the execution-context path.

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()


# ── P0-5 / C1: user.status re-verification ────────────────────────────────


async def test_assert_user_active_passes_for_active_user() -> None:
    """The active path must not raise — the run is allowed to execute."""
    executor, _, _ = _make_executor()
    session = MagicMock()
    session.scalar = AsyncMock(return_value=("active", 0))

    await executor._assert_user_active(session, "01USR0000000000000000000A")


async def test_assert_user_active_blocks_inactive_user() -> None:
    """A user whose status flipped to ``disabled`` after the run was
    queued must be rejected with a USER_DISABLED error_code-shaped prefix."""
    executor, _, _ = _make_executor()
    session = MagicMock()
    session.scalar = AsyncMock(return_value=("disabled", 0))

    with pytest.raises(ValueError, match="USER_DISABLED"):
        await executor._assert_user_active(session, "01USR0000000000000000000A")


async def test_assert_user_active_blocks_soft_deleted_user() -> None:
    """``is_deleted=1`` (soft delete) must also block execution."""
    executor, _, _ = _make_executor()
    session = MagicMock()
    session.scalar = AsyncMock(return_value=("active", 1))

    with pytest.raises(ValueError, match="USER_DISABLED"):
        await executor._assert_user_active(session, "01USR0000000000000000000A")


async def test_assert_user_active_blocks_missing_user() -> None:
    """Hard-deleted user (row gone) — also blocked, same prefix."""
    executor, _, _ = _make_executor()
    session = MagicMock()
    session.scalar = AsyncMock(return_value=None)

    with pytest.raises(ValueError, match="USER_DISABLED"):
        await executor._assert_user_active(session, "01USR0000000000000000000A")