iron/tests/test_placement_jobs.py

from sqlalchemy import select

import pytest

from app.models.entities import Blob, BlobReplica, FileEntry, FileVersion


@pytest.mark.anyio
async def test_reconcile_uses_preferred_and_excluded_backends(api_client, session_factory, fake_s3) -> None:
    stable_a = await api_client.post(
        "/api/backends",
        json={
            "name": "s3-stable-a",
            "type": "s3",
            "stability_class": "stable",
            "read_priority": 80,
            "write_priority": 80,
            "config": {"bucket": "stable-a"},
        },
    )
    stable_b = await api_client.post(
        "/api/backends",
        json={
            "name": "s3-stable-b",
            "type": "s3",
            "stability_class": "stable",
            "read_priority": 90,
            "write_priority": 90,
            "config": {"bucket": "stable-b"},
        },
    )
    opportunistic = await api_client.post(
        "/api/backends",
        json={
            "name": "s3-opp",
            "type": "s3",
            "stability_class": "opportunistic",
            "read_priority": 70,
            "write_priority": 70,
            "config": {"bucket": "opp"},
        },
    )
    stable_a_id = stable_a.json()["backend"]["id"]
    stable_b_id = stable_b.json()["backend"]["id"]
    opportunistic_id = opportunistic.json()["backend"]["id"]

    await api_client.put(
        "/api/policies/placement/image",
        json={
            "require_local": True,
            "stable_replica_count": 1,
            "opportunistic_replica_count": 1,
            "preferred_backend_ids": [stable_b_id, opportunistic_id],
            "excluded_backend_ids": [stable_a_id],
            "max_non_local_size_bytes": None,
        },
    )

    create_response = await api_client.post(
        "/files",
        headers={"Tus-Resumable": "1.0.0", "Upload-Length": "9"},
    )
    upload_id = create_response.headers["Location"].rsplit("/", 1)[-1]
    await api_client.patch(
        f"/files/{upload_id}",
        headers={
            "Tus-Resumable": "1.0.0",
            "Upload-Offset": "0",
            "Content-Type": "application/offset+octet-stream",
        },
        content=b"image-bin",
    )
    finalize_response = await api_client.post(
        f"/api/uploads/{upload_id}/finalize",
        json={"directory_id": "dir_root", "filename": "poster.jpg"},
    )
    assert finalize_response.status_code == 200

    run_response = await api_client.post("/api/jobs/run-pending")
    assert run_response.status_code == 200
    assert run_response.json()["completed"] >= 1

    file_id = finalize_response.json()["file"]["id"]
    async with session_factory() as session:
        blob_replicas = (
            await session.execute(
                select(BlobReplica.backend_id)
                .join(Blob, BlobReplica.blob_id == Blob.id)
                .join(FileVersion, Blob.file_version_id == FileVersion.id)
                .join(FileEntry, FileVersion.id == FileEntry.current_version_id)
                .where(FileEntry.id == file_id, BlobReplica.status == "ready")
            )
        ).scalars().all()
    assert "bkd_local_default" in blob_replicas
    assert stable_b_id in blob_replicas
    assert opportunistic_id in blob_replicas
    assert stable_a_id not in blob_replicas

    download_response = await api_client.get(f"/api/files/{file_id}/download")
    assert download_response.status_code == 200
    assert download_response.content == b"image-bin"


@pytest.mark.anyio
async def test_large_file_policy_can_skip_non_local_replicas(api_client, session_factory, fake_s3) -> None:
    stable_backend = await api_client.post(
        "/api/backends",
        json={
            "name": "s3-stable-a",
            "type": "s3",
            "stability_class": "stable",
            "read_priority": 80,
            "write_priority": 80,
            "config": {"bucket": "stable-a"},
        },
    )
    stable_backend_id = stable_backend.json()["backend"]["id"]

    await api_client.put(
        "/api/policies/placement/video",
        json={
            "require_local": True,
            "stable_replica_count": 1,
            "opportunistic_replica_count": 0,
            "preferred_backend_ids": [stable_backend_id],
            "excluded_backend_ids": [],
            "max_non_local_size_bytes": 4,
        },
    )

    create_response = await api_client.post(
        "/files",
        headers={"Tus-Resumable": "1.0.0", "Upload-Length": "5"},
    )
    upload_id = create_response.headers["Location"].rsplit("/", 1)[-1]
    await api_client.patch(
        f"/files/{upload_id}",
        headers={
            "Tus-Resumable": "1.0.0",
            "Upload-Offset": "0",
            "Content-Type": "application/offset+octet-stream",
        },
        content=b"hello",
    )
    finalize_response = await api_client.post(
        f"/api/uploads/{upload_id}/finalize",
        json={"directory_id": "dir_root", "filename": "clip.mp4"},
    )
    assert finalize_response.status_code == 200

    run_response = await api_client.post("/api/jobs/run-pending")
    assert run_response.status_code == 200

    async with session_factory() as session:
        blob_replicas = (
            await session.execute(
                select(BlobReplica.backend_id)
                .join(Blob, BlobReplica.blob_id == Blob.id)
                .join(FileVersion, Blob.file_version_id == FileVersion.id)
                .join(FileEntry, FileVersion.id == FileEntry.current_version_id)
                .where(
                    FileEntry.id == finalize_response.json()["file"]["id"],
                    BlobReplica.status == "ready",
                )
            )
        ).scalars().all()
    assert blob_replicas == ["bkd_local_default"]