src/backend/tests/unit/api/v1/test_files.py

import asyncio
import re
import shutil
import tempfile
from contextlib import suppress
from io import BytesIO
from pathlib import Path
from unittest.mock import MagicMock

import pytest
from asgi_lifespan import LifespanManager
from httpx import ASGITransport, AsyncClient
from langflow.services.deps import get_storage_service
from langflow.services.storage.service import StorageService
from sqlmodel import Session


@pytest.fixture
def mock_storage_service():
    # Create a mock instance of StorageService
    service = MagicMock(spec=StorageService)
    # Setup mock behaviors for the service methods as needed
    service.save_file.return_value = None
    service.get_file.return_value = b"file content"  # Binary content for files
    service.list_files.return_value = ["file1.txt", "file2.jpg"]
    service.delete_file.return_value = None

    # Mock the settings service with proper max_file_size_upload attribute
    settings_mock = MagicMock()
    settings_mock.settings = MagicMock()
    settings_mock.settings.max_file_size_upload = 1  # Default 1MB limit
    service.settings_service = settings_mock

    return service


@pytest.fixture(name="files_client")
async def files_client_fixture(
    session: Session,  # noqa: ARG001
    monkeypatch,
    request,
    load_flows_dir,
    mock_storage_service,
):
    # Set the database url to a test database
    if "noclient" in request.keywords:
        yield
    else:

        def init_app():
            db_dir = tempfile.mkdtemp()
            db_path = Path(db_dir) / "test.db"
            monkeypatch.setenv("LANGFLOW_DATABASE_URL", f"sqlite:///{db_path}")
            monkeypatch.setenv("LANGFLOW_AUTO_LOGIN", "false")
            if "load_flows" in request.keywords:
                shutil.copyfile(
                    pytest.BASIC_EXAMPLE_PATH, Path(load_flows_dir) / "c54f9130-f2fa-4a3e-b22a-3856d946351b.json"
                )
                monkeypatch.setenv("LANGFLOW_LOAD_FLOWS_PATH", load_flows_dir)
                monkeypatch.setenv("LANGFLOW_AUTO_LOGIN", "true")

            from langflow.main import create_app

            app = create_app()
            return app, db_path

        app, db_path = await asyncio.to_thread(init_app)

        app.dependency_overrides[get_storage_service] = lambda: mock_storage_service
        async with (
            LifespanManager(app, startup_timeout=None, shutdown_timeout=None) as manager,
            AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://testserver/") as client,
        ):
            yield client
        # app.dependency_overrides.clear()
        monkeypatch.undo()
        # clear the temp db
        with suppress(FileNotFoundError):
            db_path.unlink()


@pytest.fixture
async def max_file_size_upload_fixture(monkeypatch):
    monkeypatch.setenv("LANGFLOW_MAX_FILE_SIZE_UPLOAD", "1")
    yield
    monkeypatch.undo()


@pytest.fixture
async def max_file_size_upload_10mb_fixture(monkeypatch):
    monkeypatch.setenv("LANGFLOW_MAX_FILE_SIZE_UPLOAD", "10")
    yield
    monkeypatch.undo()


async def test_upload_file(files_client, created_api_key, flow):
    headers = {"x-api-key": created_api_key.api_key}

    response = await files_client.post(
        f"api/v1/files/upload/{flow.id}",
        files={"file": ("test.txt", b"test content")},
        headers=headers,
    )
    assert response.status_code == 201

    response_json = response.json()
    assert response_json["flowId"] == str(flow.id)

    # Check that the file_path matches the expected pattern
    file_path_pattern = re.compile(rf"{flow.id}/\d{{4}}-\d{{2}}-\d{{2}}_\d{{2}}-\d{{2}}-\d{{2}}_test\.txt")
    assert file_path_pattern.match(response_json["file_path"])


async def test_download_file(files_client, created_api_key, flow):
    headers = {"x-api-key": created_api_key.api_key}
    response = await files_client.get(f"api/v1/files/download/{flow.id}/test.txt", headers=headers)
    assert response.status_code == 200
    assert response.content == b"file content"


async def test_list_files(files_client, created_api_key, flow):
    headers = {"x-api-key": created_api_key.api_key}
    response = await files_client.get(f"api/v1/files/list/{flow.id}", headers=headers)
    assert response.status_code == 200
    assert response.json() == {"files": ["file1.txt", "file2.jpg"]}


async def test_delete_file(files_client, created_api_key, flow):
    headers = {"x-api-key": created_api_key.api_key}

    response = await files_client.delete(f"api/v1/files/delete/{flow.id}/test.txt", headers=headers)
    assert response.status_code == 200
    assert response.json() == {"message": "File test.txt deleted successfully"}


async def test_file_operations(client, created_api_key, flow):
    headers = {"x-api-key": created_api_key.api_key}
    flow_id = flow.id
    file_name = "test.txt"
    file_content = b"Hello, world!"

    # Step 1: Upload the file
    response = await client.post(
        f"api/v1/files/upload/{flow_id}",
        files={"file": (file_name, file_content)},
        headers=headers,
    )
    assert response.status_code == 201

    response_json = response.json()
    assert response_json["flowId"] == str(flow_id)

    # Check that the file_path matches the expected pattern
    file_path_pattern = re.compile(rf"{flow_id}/\d{{4}}-\d{{2}}-\d{{2}}_\d{{2}}-\d{{2}}-\d{{2}}_{file_name}")
    assert file_path_pattern.match(response_json["file_path"])

    # Extract the full file name with timestamp from the response
    full_file_name = response_json["file_path"].split("/")[-1]

    # Step 2: List files in the folder
    response = await client.get(f"api/v1/files/list/{flow_id}", headers=headers)
    assert response.status_code == 200
    assert full_file_name in response.json()["files"]

    # Step 3: Download the file and verify its content
    response = await client.get(f"api/v1/files/download/{flow_id}/{full_file_name}", headers=headers)
    assert response.status_code == 200
    assert response.content == file_content
    assert response.headers["content-type"] == "application/octet-stream"

    # Step 4: Delete the file
    response = await client.delete(f"api/v1/files/delete/{flow_id}/{full_file_name}", headers=headers)
    assert response.status_code == 200
    assert response.json() == {"message": f"File {full_file_name} deleted successfully"}

    # Verify that the file is indeed deleted
    response = await client.get(f"api/v1/files/list/{flow_id}", headers=headers)
    assert full_file_name not in response.json()["files"]


@pytest.mark.usefixtures("max_file_size_upload_fixture")
async def test_upload_file_size_limit(files_client, created_api_key, flow):
    headers = {"x-api-key": created_api_key.api_key}

    # Test file under the limit (500KB)
    small_content = b"x" * (500 * 1024)
    small_file = ("small_file.txt", small_content, "application/octet-stream")
    headers["Content-Length"] = str(len(small_content))
    response = await files_client.post(
        f"api/v1/files/upload/{flow.id}",
        files={"file": small_file},
        headers=headers,
    )
    assert response.status_code == 201, f"Expected 201, got {response.status_code}: {response.json()}"

    # Test file over the limit (1MB + 1KB)
    large_content = b"x" * (1024 * 1024 + 1024)

    bio = BytesIO(large_content)
    headers["Content-Length"] = str(len(large_content))
    response = await files_client.post(
        f"api/v1/files/upload/{flow.id}",
        files={"file": ("large_file.txt", bio, "application/octet-stream")},
        headers=headers,
    )

    assert response.status_code == 413, f"Expected 413, got {response.status_code}: {response.json()}"
    assert "Content size limit exceeded. Maximum allowed is 1MB and got 1.001MB." in response.json()["detail"]