yunke/yunke-server/tests/test_api_media.py

"""Tests for GET /api/media/{media_hash} — media serve API (Task 2.8).

Drives the FastAPI app over an in-process httpx transport (``ASGITransport``) —
no real network socket. The ``get_session`` dependency is overridden to an
isolated in-memory async SQLite session (``StaticPool`` so every request sees
the same seeded schema/rows). Media blobs are written to a real ``tmp_path``
file and referenced by absolute path from the seeded ``Media`` row.

Coverage:
* ``GET /api/media/{hash}`` for a seeded row -> 200, correct ``Content-Type``,
  and the exact bytes streamed back.
* Unknown hash -> 404.
* Row exists but the underlying file is missing on disk -> 404.
* HTTP Range request -> ``206`` with the requested partial bytes (FileResponse
  handles Range), proving the endpoint supports range / partial-content.
* A row whose ``mime_type`` is NULL falls back to an ext-derived content type.
"""

from __future__ import annotations

from pathlib import Path

import pytest
from httpx import ASGITransport, AsyncClient
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.pool import StaticPool

from app.api import deps
from app.api.main import app
from app.config import get_settings
from app.db.base import Base
from app.db.models import Media

PNG_BYTES = b"\x89PNG\r\n\x1a\n fake image body"


@pytest.fixture
async def client(tmp_path: Path, monkeypatch):
    """Yield an ``AsyncClient`` against ``app`` with a seeded in-memory DB.

    Seeds one ``Media`` row (``abc123``) pointing at a real PNG-bytes file
    under ``tmp_path``, plus a second row (``nopng``) whose ``local_path`` is
    deliberately missing on disk (to exercise the file-missing 404 path), plus
    a third (``nomime``) whose ``mime_type`` is NULL (to exercise the
    ext-derived content-type fallback).

    ``MEDIA_DIR`` is pointed at ``tmp_path`` so the endpoint's path-traversal
    sandbox check (file must live under ``Settings.media_dir``) admits the
    seeded rows. The settings cache is cleared around the override so the new
    value is picked up.
    """
    monkeypatch.setenv("MEDIA_DIR", str(tmp_path))
    get_settings.cache_clear()

    engine = create_async_engine(
        "sqlite+aiosqlite:///:memory:",
        poolclass=StaticPool,
        connect_args={"check_same_thread": False},
    )
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

    media_file = tmp_path / "abc123.png"
    media_file.write_bytes(PNG_BYTES)

    # A second file with no mime_type recorded, to test ext-derived fallback.
    fallback_file = tmp_path / "fallback.webp"
    fallback_file.write_bytes(b"WEBPBYTES")

    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
    async with SessionLocal() as session:
        session.add_all(
            [
                Media(
                    media_hash="abc123",
                    channel="wechat",
                    media_type="image",
                    original_url="https://wx.oss.yunkecn.com/x/abc.png",
                    local_path=str(media_file),
                    size_bytes=len(PNG_BYTES),
                    mime_type="image/png",
                ),
                Media(
                    media_hash="nopng",
                    channel="wechat",
                    media_type="image",
                    local_path=str(tmp_path / "does-not-exist.png"),
                    size_bytes=1,
                    mime_type="image/png",
                ),
                Media(
                    media_hash="nomime",
                    channel="wechat",
                    media_type="image",
                    local_path=str(fallback_file),
                    size_bytes=len(b"WEBPBYTES"),
                    mime_type=None,
                ),
            ]
        )
        await session.commit()

    async def _override_session():
        async with SessionLocal() as s:
            yield s

    app.dependency_overrides[deps.get_session] = _override_session
    try:
        async with AsyncClient(
            transport=ASGITransport(app=app), base_url="http://test"
        ) as ac:
            yield ac
    finally:
        app.dependency_overrides.clear()
        await engine.dispose()
        get_settings.cache_clear()


@pytest.fixture
async def escape_client(tmp_path: Path, monkeypatch):
    """Variant of :func:`client` that pins ``MEDIA_DIR`` to a sandbox subdir
    while seeding a row whose ``local_path`` points OUTSIDE that sandbox.

    Used to verify the path-traversal defense: an outside-sandbox file must be
    refused with 404 even when it exists on disk and the row is otherwise valid.
    """
    sandbox = tmp_path / "media_dir"
    sandbox.mkdir()
    outside = tmp_path / "outside"
    outside.mkdir()

    monkeypatch.setenv("MEDIA_DIR", str(sandbox))
    get_settings.cache_clear()

    engine = create_async_engine(
        "sqlite+aiosqlite:///:memory:",
        poolclass=StaticPool,
        connect_args={"check_same_thread": False},
    )
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

    escape_file = outside / "secret.png"
    escape_file.write_bytes(b"TOPSECRET")

    SessionLocal = async_sessionmaker(engine, expire_on_commit=False)
    async with SessionLocal() as session:
        session.add(
            Media(
                media_hash="escape",
                channel="wechat",
                media_type="image",
                original_url="https://wx.oss.yunkecn.com/x/secret.png",
                local_path=str(escape_file),
                size_bytes=len(b"TOPSECRET"),
                mime_type="image/png",
            )
        )
        await session.commit()

    async def _override_session():
        async with SessionLocal() as s:
            yield s

    app.dependency_overrides[deps.get_session] = _override_session
    try:
        async with AsyncClient(
            transport=ASGITransport(app=app), base_url="http://test"
        ) as ac:
            yield ac
    finally:
        app.dependency_overrides.clear()
        await engine.dispose()
        get_settings.cache_clear()


async def test_get_media_returns_bytes_with_content_type(client: AsyncClient):
    """Seeded hash -> 200, image/png content-type, exact bytes streamed back."""
    r = await client.get("/api/media/abc123")
    assert r.status_code == 200
    assert r.headers["content-type"].startswith("image/png")
    assert r.content == PNG_BYTES


async def test_get_media_unknown_hash_returns_404(client: AsyncClient):
    """An unknown media_hash -> 404."""
    r = await client.get("/api/media/does-not-exist")
    assert r.status_code == 404


async def test_get_media_missing_file_on_disk_returns_404(client: AsyncClient):
    """Row present but file absent from disk -> 404 (defensive)."""
    r = await client.get("/api/media/nopng")
    assert r.status_code == 404


async def test_get_media_supports_http_range(client: AsyncClient):
    """Range request -> 206 Partial Content with the requested byte slice.

    Starlette's ``FileResponse`` honors ``Range``; this pins the behavior so a
    future swap to a hand-rolled StreamingResponse would have to keep it.
    """
    r = await client.get("/api/media/abc123", headers={"Range": "bytes=0-3"})
    assert r.status_code == 206
    assert r.content == PNG_BYTES[:4]
    assert r.headers["content-length"] == "4"
    assert r.headers["content-range"].startswith("bytes 0-3/")


async def test_get_media_falls_back_to_ext_when_mime_null(client: AsyncClient):
    """mime_type NULL on the row -> content-type guessed from the .webp ext."""
    r = await client.get("/api/media/nomime")
    assert r.status_code == 200
    assert r.headers["content-type"].startswith("image/webp")
    assert r.content == b"WEBPBYTES"


async def test_get_media_refuses_path_outside_media_dir(escape_client: AsyncClient):
    """A row whose local_path escapes media_dir is refused with 404 — even
    though the file exists on disk and the row is otherwise valid.

    Pins the path-traversal defense: a hostile or buggy ``local_path`` value
    must never let the endpoint serve files outside the configured sandbox.
    """
    r = await escape_client.get("/api/media/escape")
    assert r.status_code == 404
    # The file must NOT be served — only the 404 JSON body is returned.
    assert b"TOPSECRET" not in r.content
    assert "media not found" in r.json()["detail"]