prefecthq--fastmcp/tests/apps/test_file_upload.py

"""Tests for the FileUpload provider."""

import base64

import pytest

from fastmcp import FastMCP
from fastmcp.apps.file_upload import FileUpload, _b64_decoded_size
from fastmcp.server.providers.addressing import hashed_backend_name


class TestB64DecodedSize:
    """Unit tests for the _b64_decoded_size helper."""

    @pytest.mark.parametrize("size", [0, 1, 2, 3, 4, 50, 99, 100, 101, 1000])
    def test_matches_actual_decode(self, size: int):
        data = b"x" * size
        b64 = base64.b64encode(data).decode()
        assert _b64_decoded_size(b64) == size

    def test_empty_string(self):
        assert _b64_decoded_size("") == 0

    def test_no_padding(self):
        # 3 bytes → 4 base64 chars, no padding
        assert _b64_decoded_size(base64.b64encode(b"abc").decode()) == 3

    def test_one_pad(self):
        # 2 bytes → 4 base64 chars with 1 '='
        assert _b64_decoded_size(base64.b64encode(b"ab").decode()) == 2

    def test_two_pads(self):
        # 1 byte → 4 base64 chars with 2 '='
        assert _b64_decoded_size(base64.b64encode(b"a").decode()) == 1


def _make_file(
    name: str = "test.txt",
    content: str = "hello world",
    mime_type: str = "text/plain",
) -> dict:
    data = base64.b64encode(content.encode()).decode()
    return {
        "name": name,
        "size": len(content),
        "type": mime_type,
        "data": data,
    }


class TestFileUploadProvider:
    async def test_basic_store_and_list(self):
        server = FastMCP("test", providers=[FileUpload()])
        files = [_make_file()]

        result = await server.call_tool(
            hashed_backend_name("Files", "store_files"), {"files": files}
        )
        text = result.content[0].text  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
        assert "test.txt" in text

        result = await server.call_tool("list_files", {})
        text = result.content[0].text  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
        assert "test.txt" in text

    async def test_read_text_file(self):
        server = FastMCP("test", providers=[FileUpload()])
        files = [_make_file(content="DON'T PANIC")]

        await server.call_tool(
            hashed_backend_name("Files", "store_files"), {"files": files}
        )

        result = await server.call_tool("read_file", {"name": "test.txt"})
        text = result.content[0].text  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
        assert "DON'T PANIC" in text

    async def test_read_binary_file(self):
        server = FastMCP("test", providers=[FileUpload()])
        data = base64.b64encode(b"\x00\x01\x02\xff").decode()
        files = [{"name": "image.png", "size": 4, "type": "image/png", "data": data}]

        await server.call_tool(
            hashed_backend_name("Files", "store_files"), {"files": files}
        )

        result = await server.call_tool("read_file", {"name": "image.png"})
        text = result.content[0].text  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
        assert "content_base64" in text

    async def test_read_missing_file_raises(self):
        server = FastMCP("test", providers=[FileUpload()])

        with pytest.raises(Exception, match="not found"):
            await server.call_tool("read_file", {"name": "nope.txt"})

    async def test_multiple_files(self):
        server = FastMCP("test", providers=[FileUpload()])
        files = [
            _make_file("a.txt", "aaa"),
            _make_file("b.txt", "bbb"),
        ]

        await server.call_tool(
            hashed_backend_name("Files", "store_files"), {"files": files}
        )

        result = await server.call_tool("list_files", {})
        text = result.content[0].text  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
        assert "a.txt" in text
        assert "b.txt" in text

    async def test_overwrite_file(self):
        server = FastMCP("test", providers=[FileUpload()])

        await server.call_tool(
            hashed_backend_name("Files", "store_files"),
            {"files": [_make_file(content="version 1")]},
        )
        await server.call_tool(
            hashed_backend_name("Files", "store_files"),
            {"files": [_make_file(content="version 2")]},
        )

        result = await server.call_tool("read_file", {"name": "test.txt"})
        text = result.content[0].text  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
        assert "version 2" in text

    async def test_custom_name(self):
        server = FastMCP("test", providers=[FileUpload(name="Uploads")])

        tools = await server.list_tools()
        tool_names = [t.name for t in tools]
        assert "file_manager" in tool_names

        # The hash uses the app's actual name ("Uploads"), not the default.
        files = [_make_file()]
        result = await server.call_tool(
            hashed_backend_name("Uploads", "store_files"), {"files": files}
        )
        text = result.content[0].text  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
        assert "test.txt" in text

    async def test_ui_tool_visible_backend_hidden(self):
        server = FastMCP("test", providers=[FileUpload()])

        tools = await server.list_tools()
        tool_names = [t.name for t in tools]

        assert "file_manager" in tool_names
        assert "list_files" in tool_names
        assert "read_file" in tool_names
        assert "store_files" not in tool_names

    async def test_max_file_size_enforced_server_side(self):
        server = FastMCP("test", providers=[FileUpload(max_file_size=100)])
        big_file = _make_file(content="x" * 200)

        with pytest.raises(Exception, match="exceeds max size"):
            await server.call_tool(
                hashed_backend_name("Files", "store_files"), {"files": [big_file]}
            )

    async def test_max_file_size_checks_actual_data_not_reported_size(self):
        """Size limit should be enforced on actual base64 payload, not the
        client-reported ``size`` field which can be spoofed."""
        server = FastMCP("test", providers=[FileUpload(max_file_size=100)])

        big_content = "x" * 200
        big_b64 = base64.b64encode(big_content.encode()).decode()
        spoofed_file = {
            "name": "spoofed.bin",
            "size": 1,  # lies about size
            "type": "application/octet-stream",
            "data": big_b64,
        }

        with pytest.raises(Exception, match="exceeds max size"):
            await server.call_tool(
                hashed_backend_name("Files", "store_files"), {"files": [spoofed_file]}
            )


class TestFileUploadSubclass:
    async def test_custom_storage(self):
        """Subclassing lets users provide their own persistence."""
        stored: dict[str, dict] = {}

        class MemoryUpload(FileUpload):
            def on_store(self, files: list[dict], ctx) -> list[dict]:
                for f in files:
                    stored[f["name"]] = f
                return [
                    {
                        "name": f["name"],
                        "type": f["type"],
                        "size": f["size"],
                        "size_display": "?",
                        "uploaded_at": "now",
                    }
                    for f in files
                ]

            def on_list(self, ctx) -> list[dict]:
                return [
                    {
                        "name": f["name"],
                        "type": f["type"],
                        "size": f["size"],
                        "size_display": "?",
                        "uploaded_at": "now",
                    }
                    for f in stored.values()
                ]

            def on_read(self, name: str, ctx) -> dict:
                if name not in stored:
                    raise ValueError(f"Not found: {name}")
                f = stored[name]
                return {"name": f["name"], "content": "custom read"}

        server = FastMCP("test", providers=[MemoryUpload()])
        files = [_make_file()]

        await server.call_tool(
            hashed_backend_name("Files", "store_files"), {"files": files}
        )

        assert "test.txt" in stored

        result = await server.call_tool("read_file", {"name": "test.txt"})
        text = result.content[0].text  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
        assert "custom read" in text