fix: updates file size limit to use middleware and add tests for uploads (#4883)

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

@@ -0,0 +1,208 @@
+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"]