perf: optimize url upload with async i/o

Replaced synchronous `requests.get` and `open().write` in the `process_url` endpoint with `httpx.AsyncClient` and `aiofiles.open`. This prevents the FastAPI event loop from blocking during large file downloads.

Updated test suite in `tests/test_url_upload.py` to use `AsyncMock` to mock `httpx.AsyncClient.stream` contexts and async generators properly, covering all original conditions and HTTP error handling paths.

Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
google-labs-jules[bot]
2026-03-16 09:41:17 +00:00
parent 2732dafba9
commit 320a2acedd
4 changed files with 388 additions and 161 deletions
+56 -56
View File
@@ -9,7 +9,8 @@ import urllib.parse
import uuid import uuid
from typing import Optional from typing import Optional
import requests import aiofiles
import httpx
from fastapi import APIRouter, HTTPException, Request from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel, HttpUrl, field_validator from pydantic import BaseModel, HttpUrl, field_validator
@@ -153,67 +154,66 @@ async def process_url(request: Request, url_request: URLUploadRequest):
logger.info(f"Downloading file from URL: {url}") logger.info(f"Downloading file from URL: {url}")
# Use configured timeout to prevent hanging # Use configured timeout to prevent hanging
response = requests.get( async with httpx.AsyncClient(
url,
timeout=settings.http_request_timeout, timeout=settings.http_request_timeout,
stream=True, # Stream to handle large files follow_redirects=True,
allow_redirects=True, # Follow redirects
headers={ headers={
"User-Agent": "DocuElevate/1.0", # Identify ourselves "User-Agent": "DocuElevate/1.0", # Identify ourselves
}, },
) ) as client:
response.raise_for_status() async with client.stream("GET", url) as response:
response.raise_for_status()
# Validate content type # Validate content type
content_type = response.headers.get("Content-Type", "") content_type = response.headers.get("Content-Type", "")
if not validate_file_type(content_type, safe_filename): if not validate_file_type(content_type, safe_filename):
raise HTTPException( raise HTTPException(
status_code=400, status_code=400,
detail=f"Unsupported file type: {content_type}. " detail=f"Unsupported file type: {content_type}. "
"Supported types: PDF, Office documents, images, plain text", "Supported types: PDF, Office documents, images, plain text",
) )
# Check content length before downloading # Check content length before downloading
content_length = response.headers.get("Content-Length") content_length = response.headers.get("Content-Length")
if content_length: if content_length:
file_size = int(content_length) file_size = int(content_length)
max_size = settings.max_upload_size max_size = settings.max_upload_size
if file_size > max_size: if file_size > max_size:
raise HTTPException(
status_code=413,
detail=f"File too large: {file_size} bytes (max {max_size} bytes)",
)
# Generate unique filename
unique_id = str(uuid.uuid4())
if "." in safe_filename:
file_extension = safe_filename.rsplit(".", 1)[1]
target_filename = f"{unique_id}.{file_extension}"
else:
target_filename = unique_id
target_path = os.path.join(settings.workdir, target_filename)
# Download file in chunks to handle large files
downloaded_size = 0
max_size = settings.max_upload_size
with open(target_path, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
downloaded_size += len(chunk)
# Check size during download
if downloaded_size > max_size:
# Remove partial file
f.close()
os.remove(target_path)
raise HTTPException( raise HTTPException(
status_code=413, status_code=413,
detail=f"File too large: exceeded {max_size} bytes during download", detail=f"File too large: {file_size} bytes (max {max_size} bytes)",
) )
# Generate unique filename
unique_id = str(uuid.uuid4())
if "." in safe_filename:
file_extension = safe_filename.rsplit(".", 1)[1]
target_filename = f"{unique_id}.{file_extension}"
else:
target_filename = unique_id
target_path = os.path.join(settings.workdir, target_filename)
# Download file in chunks to handle large files
downloaded_size = 0
max_size = settings.max_upload_size
async with aiofiles.open(target_path, "wb") as f:
async for chunk in response.aiter_bytes(chunk_size=8192):
if chunk:
await f.write(chunk)
downloaded_size += len(chunk)
# Check size during download
if downloaded_size > max_size:
# Remove partial file
await f.close()
os.remove(target_path)
raise HTTPException(
status_code=413,
detail=f"File too large: exceeded {max_size} bytes during download",
)
logger.info(f"Downloaded file from URL '{url}' as '{target_filename}' ({downloaded_size} bytes)") logger.info(f"Downloaded file from URL '{url}' as '{target_filename}' ({downloaded_size} bytes)")
# Enqueue for processing # Enqueue for processing
@@ -227,19 +227,19 @@ async def process_url(request: Request, url_request: URLUploadRequest):
"size": downloaded_size, "size": downloaded_size,
} }
except requests.exceptions.Timeout: except httpx.TimeoutException:
logger.error(f"Timeout while downloading file from URL: {url}") logger.error(f"Timeout while downloading file from URL: {url}")
raise HTTPException(status_code=408, detail="Request timeout: server took too long to respond") raise HTTPException(status_code=408, detail="Request timeout: server took too long to respond")
except requests.exceptions.ConnectionError as e: except httpx.ConnectError as e:
logger.error(f"Connection error while downloading file from URL: {url} - {str(e)}") logger.error(f"Connection error while downloading file from URL: {url} - {str(e)}")
raise HTTPException(status_code=502, detail=f"Failed to connect to URL: {str(e)}") raise HTTPException(status_code=502, detail=f"Failed to connect to URL: {str(e)}")
except requests.exceptions.HTTPError as e: except httpx.HTTPStatusError as e:
logger.error(f"HTTP error while downloading file from URL: {url} - {str(e)}") logger.error(f"HTTP error while downloading file from URL: {url} - {str(e)}")
raise HTTPException(status_code=e.response.status_code, detail=f"HTTP error: {str(e)}") raise HTTPException(status_code=e.response.status_code, detail=f"HTTP error: {str(e)}")
except requests.exceptions.RequestException as e: except httpx.RequestError as e:
logger.error(f"Error downloading file from URL: {url} - {str(e)}") logger.error(f"Error downloading file from URL: {url} - {str(e)}")
raise HTTPException(status_code=500, detail=f"Failed to download file: {str(e)}") raise HTTPException(status_code=500, detail=f"Failed to download file: {str(e)}")
+49
View File
@@ -0,0 +1,49 @@
import asyncio
import time
from unittest.mock import Mock, patch
from app.api.url_upload import process_url, URLUploadRequest
from app.config import settings
async def main():
# Mock request and URLUploadRequest
request = Mock()
url_request = URLUploadRequest(url="https://example.com/file.pdf")
# Generate a large chunk
large_chunk = b"A" * 8192
num_chunks = 10000 # 8192 * 10000 = ~80MB
mock_response = Mock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf"}
mock_response.iter_content = Mock(return_value=[large_chunk] * num_chunks)
# For async client later
class AsyncMockResponse:
def __init__(self):
self.status_code = 200
self.headers = {"Content-Type": "application/pdf"}
def raise_for_status(self):
pass
async def aiter_bytes(self, chunk_size):
for _ in range(num_chunks):
yield large_chunk
async_mock_response = AsyncMockResponse()
# We will mock requests.get for synchronous, httpx.AsyncClient.get for asynchronous
# Test sync
start_time = time.time()
with patch("app.api.url_upload.requests.get", return_value=mock_response), \
patch("app.api.url_upload.process_document"):
try:
await process_url(request=request, url_request=url_request)
except Exception as e:
print(f"Error: {e}")
end_time = time.time()
print(f"Original execution time (sync writing): {end_time - start_time:.4f} seconds")
if __name__ == "__main__":
asyncio.run(main())
+75
View File
@@ -0,0 +1,75 @@
import asyncio
import time
import os
import shutil
import tempfile
from unittest.mock import Mock, patch
from fastapi import HTTPException
from app.api.url_upload import process_url, URLUploadRequest
from app.config import settings
async def main():
# Setup test dir
test_dir = tempfile.mkdtemp()
settings.workdir = test_dir
# Mock request and URLUploadRequest
request = Mock()
url_request = URLUploadRequest(url="https://example.com/file.pdf")
# Generate a large chunk
chunk_size = 8192
num_chunks = 20000 # 20000 * 8192 = ~160MB
large_chunk = b"A" * chunk_size
class SyncMockResponse:
def __init__(self):
self.status_code = 200
self.headers = {"Content-Type": "application/pdf"}
def raise_for_status(self):
pass
def iter_content(self, chunk_size):
for _ in range(num_chunks):
# sleep slightly to simulate network latency, otherwise OS file cache obscures the difference
time.sleep(0.0001)
yield large_chunk
sync_mock_response = SyncMockResponse()
class AsyncMockResponse:
def __init__(self):
self.status_code = 200
self.headers = {"Content-Type": "application/pdf"}
self.is_success = True
self.status_code = 200
def raise_for_status(self):
pass
async def aiter_bytes(self, chunk_size=8192):
for _ in range(num_chunks):
await asyncio.sleep(0.0001)
yield large_chunk
class AsyncMockContext:
async def __aenter__(self):
return AsyncMockResponse()
async def __aexit__(self, exc_type, exc_val, exc_tb):
pass
async_mock_response = AsyncMockResponse()
# Test sync
start_time = time.time()
with patch("app.api.url_upload.requests.get", return_value=sync_mock_response), \
patch("app.api.url_upload.process_document"):
try:
await process_url(request=request, url_request=url_request)
except Exception as e:
print(f"Error (sync): {e}")
end_time = time.time()
print(f"Original execution time (sync writing): {end_time - start_time:.4f} seconds")
shutil.rmtree(test_dir)
if __name__ == "__main__":
asyncio.run(main())
+208 -105
View File
@@ -2,10 +2,10 @@
Tests for URL-based file upload functionality Tests for URL-based file upload functionality
""" """
from unittest.mock import Mock, patch from unittest.mock import AsyncMock, MagicMock, Mock, patch
import httpx
import pytest import pytest
import requests
@pytest.mark.unit @pytest.mark.unit
@@ -165,17 +165,24 @@ class TestURLUploadValidation:
class TestURLUploadEndpoint: class TestURLUploadEndpoint:
"""Integration tests for URL upload endpoint""" """Integration tests for URL upload endpoint"""
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_requires_authentication(self, mock_process_document, mock_requests_get, client, monkeypatch): def test_process_url_requires_authentication(self, mock_process_document, mock_stream, client, monkeypatch):
"""Test that endpoint requires authentication when auth is enabled""" """Test that endpoint requires authentication when auth is enabled"""
# Mock successful download to prevent actual HTTP requests # Mock successful download to prevent actual HTTP requests
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "1024"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "1024"}
mock_response.iter_content = Mock(return_value=[b"PDF content"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF content"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
# Mock Celery task # Mock Celery task
mock_task = Mock() mock_task = Mock()
@@ -192,17 +199,24 @@ class TestURLUploadEndpoint:
# (like no mocking). We're just checking the endpoint exists and is reachable. # (like no mocking). We're just checking the endpoint exists and is reachable.
assert response.status_code != 404 # Endpoint should exist assert response.status_code != 404 # Endpoint should exist
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_success(self, mock_process_document, mock_requests_get, client, tmp_path): def test_process_url_success(self, mock_process_document, mock_stream, client, tmp_path):
"""Test successful URL processing""" """Test successful URL processing"""
# Mock successful download # Mock successful download
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "1024"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "1024"}
mock_response.iter_content = Mock(return_value=[b"PDF content here"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF content here"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
# Mock Celery task # Mock Celery task
mock_task = Mock() mock_task = Mock()
@@ -219,8 +233,8 @@ class TestURLUploadEndpoint:
assert "filename" in data assert "filename" in data
assert "size" in data assert "size" in data
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_blocks_private_ip(self, mock_requests_get, client): def test_process_url_blocks_private_ip(self, mock_stream, client):
"""Test that private IPs are blocked""" """Test that private IPs are blocked"""
response = client.post("/api/process-url", json={"url": "http://192.168.1.1/file.pdf"}) response = client.post("/api/process-url", json={"url": "http://192.168.1.1/file.pdf"})
@@ -229,10 +243,10 @@ class TestURLUploadEndpoint:
assert "private/internal" in data["detail"] assert "private/internal" in data["detail"]
# Should not make HTTP request # Should not make HTTP request
mock_requests_get.assert_not_called() mock_stream.assert_not_called()
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_blocks_localhost(self, mock_requests_get, client): def test_process_url_blocks_localhost(self, mock_stream, client):
"""Test that localhost is blocked""" """Test that localhost is blocked"""
response = client.post("/api/process-url", json={"url": "http://localhost/file.pdf"}) response = client.post("/api/process-url", json={"url": "http://localhost/file.pdf"})
@@ -241,10 +255,10 @@ class TestURLUploadEndpoint:
assert "private/internal" in data["detail"] assert "private/internal" in data["detail"]
# Should not make HTTP request # Should not make HTTP request
mock_requests_get.assert_not_called() mock_stream.assert_not_called()
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_blocks_metadata_endpoint(self, mock_requests_get, client): def test_process_url_blocks_metadata_endpoint(self, mock_stream, client):
"""Test that cloud metadata endpoints are blocked""" """Test that cloud metadata endpoints are blocked"""
response = client.post("/api/process-url", json={"url": "http://169.254.169.254/latest/meta-data/"}) response = client.post("/api/process-url", json={"url": "http://169.254.169.254/latest/meta-data/"})
@@ -254,17 +268,20 @@ class TestURLUploadEndpoint:
assert "metadata" in data["detail"] or "private" in data["detail"] assert "metadata" in data["detail"] or "private" in data["detail"]
# Should not make HTTP request # Should not make HTTP request
mock_requests_get.assert_not_called() mock_stream.assert_not_called()
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_invalid_file_type(self, mock_requests_get, client): def test_process_url_invalid_file_type(self, mock_stream, client):
"""Test that invalid file types are rejected""" """Test that invalid file types are rejected"""
# Mock response with executable content-type # Mock response with executable content-type
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/x-executable"} mock_response.headers = {"Content-Type": "application/x-executable"}
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
response = client.post("/api/process-url", json={"url": "https://example.com/malware.exe"}) response = client.post("/api/process-url", json={"url": "https://example.com/malware.exe"})
@@ -272,21 +289,24 @@ class TestURLUploadEndpoint:
data = response.json() data = response.json()
assert "Unsupported file type" in data["detail"] assert "Unsupported file type" in data["detail"]
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_file_too_large_by_header(self, mock_process_document, mock_requests_get, client): def test_process_url_file_too_large_by_header(self, mock_process_document, mock_stream, client):
"""Test that files too large are rejected based on Content-Length header""" """Test that files too large are rejected based on Content-Length header"""
from app.config import settings from app.config import settings
# Mock response with large content-length # Mock response with large content-length
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = { mock_response.headers = {
"Content-Type": "application/pdf", "Content-Type": "application/pdf",
"Content-Length": str(settings.max_upload_size + 1000), "Content-Length": str(settings.max_upload_size + 1000),
} }
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
response = client.post("/api/process-url", json={"url": "https://example.com/huge.pdf"}) response = client.post("/api/process-url", json={"url": "https://example.com/huge.pdf"})
@@ -297,10 +317,10 @@ class TestURLUploadEndpoint:
# Should not process document # Should not process document
mock_process_document.delay.assert_not_called() mock_process_document.delay.assert_not_called()
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_timeout_error(self, mock_requests_get, client): def test_process_url_timeout_error(self, mock_stream, client):
"""Test handling of timeout errors""" """Test handling of timeout errors"""
mock_requests_get.side_effect = requests.exceptions.Timeout("Request timed out") mock_stream.side_effect = httpx.TimeoutException("Request timed out")
response = client.post("/api/process-url", json={"url": "https://example.com/slow.pdf"}) response = client.post("/api/process-url", json={"url": "https://example.com/slow.pdf"})
@@ -308,10 +328,10 @@ class TestURLUploadEndpoint:
data = response.json() data = response.json()
assert "timeout" in data["detail"].lower() assert "timeout" in data["detail"].lower()
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_connection_error(self, mock_requests_get, client): def test_process_url_connection_error(self, mock_stream, client):
"""Test handling of connection errors""" """Test handling of connection errors"""
mock_requests_get.side_effect = requests.exceptions.ConnectionError("Failed to connect") mock_stream.side_effect = httpx.ConnectError("Failed to connect")
response = client.post("/api/process-url", json={"url": "https://example.com/file.pdf"}) response = client.post("/api/process-url", json={"url": "https://example.com/file.pdf"})
@@ -319,15 +339,16 @@ class TestURLUploadEndpoint:
data = response.json() data = response.json()
assert "connect" in data["detail"].lower() assert "connect" in data["detail"].lower()
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_http_error_404(self, mock_requests_get, client): def test_process_url_http_error_404(self, mock_stream, client):
"""Test handling of HTTP 404 errors""" """Test handling of HTTP 404 errors"""
mock_response = Mock() # When raising HTTPStatusError, httpx requires request and response arguments
# For our code, we just need it to hit the exception handler and check status code
mock_request = MagicMock()
mock_response = MagicMock()
mock_response.status_code = 404 mock_response.status_code = 404
mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError(
"404 Not Found", response=mock_response mock_stream.side_effect = httpx.HTTPStatusError("404 Not Found", request=mock_request, response=mock_response)
)
mock_requests_get.return_value = mock_response
response = client.post("/api/process-url", json={"url": "https://example.com/notfound.pdf"}) response = client.post("/api/process-url", json={"url": "https://example.com/notfound.pdf"})
@@ -335,17 +356,24 @@ class TestURLUploadEndpoint:
data = response.json() data = response.json()
assert "HTTP error" in data["detail"] assert "HTTP error" in data["detail"]
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_with_custom_filename(self, mock_process_document, mock_requests_get, client, tmp_path): def test_process_url_with_custom_filename(self, mock_process_document, mock_stream, client, tmp_path):
"""Test URL upload with custom filename""" """Test URL upload with custom filename"""
# Mock successful download # Mock successful download
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "1024"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "1024"}
mock_response.iter_content = Mock(return_value=[b"PDF content"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF content"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
# Mock Celery task # Mock Celery task
mock_task = Mock() mock_task = Mock()
@@ -361,17 +389,24 @@ class TestURLUploadEndpoint:
data = response.json() data = response.json()
assert data["filename"] == "my-document.pdf" assert data["filename"] == "my-document.pdf"
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_extracts_filename_from_url(self, mock_process_document, mock_requests_get, client, tmp_path): def test_process_url_extracts_filename_from_url(self, mock_process_document, mock_stream, client, tmp_path):
"""Test that filename is extracted from URL when not provided""" """Test that filename is extracted from URL when not provided"""
# Mock successful download # Mock successful download
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "1024"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "1024"}
mock_response.iter_content = Mock(return_value=[b"PDF content"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF content"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
# Mock Celery task # Mock Celery task
mock_task = Mock() mock_task = Mock()
@@ -386,9 +421,9 @@ class TestURLUploadEndpoint:
# Should extract "annual-report.pdf" from URL # Should extract "annual-report.pdf" from URL
assert "annual-report" in data["filename"] assert "annual-report" in data["filename"]
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_file_size_during_download(self, mock_process_document, mock_requests_get, client): def test_process_url_file_size_during_download(self, mock_process_document, mock_stream, client):
"""Test that file size is checked during download""" """Test that file size is checked during download"""
from app.config import settings from app.config import settings
@@ -396,12 +431,19 @@ class TestURLUploadEndpoint:
large_chunk = b"x" * (settings.max_upload_size + 1000) large_chunk = b"x" * (settings.max_upload_size + 1000)
# Mock response without Content-Length header # Mock response without Content-Length header
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf"} # No Content-Length mock_response.headers = {"Content-Type": "application/pdf"} # No Content-Length
mock_response.iter_content = Mock(return_value=[large_chunk])
async def mock_aiter_bytes(chunk_size=None):
yield large_chunk
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
response = client.post("/api/process-url", json={"url": "https://example.com/big.pdf"}) response = client.post("/api/process-url", json={"url": "https://example.com/big.pdf"})
@@ -412,10 +454,10 @@ class TestURLUploadEndpoint:
# Should not process document # Should not process document
mock_process_document.delay.assert_not_called() mock_process_document.delay.assert_not_called()
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_request_exception(self, mock_requests_get, client): def test_process_url_request_exception(self, mock_stream, client):
"""Test handling of generic RequestException""" """Test handling of generic RequestError"""
mock_requests_get.side_effect = requests.exceptions.RequestException("Generic request error") mock_stream.side_effect = httpx.RequestError("Generic request error")
response = client.post("/api/process-url", json={"url": "https://example.com/file.pdf"}) response = client.post("/api/process-url", json={"url": "https://example.com/file.pdf"})
@@ -423,16 +465,23 @@ class TestURLUploadEndpoint:
data = response.json() data = response.json()
assert "Failed to download file" in data["detail"] assert "Failed to download file" in data["detail"]
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_oserror_during_save(self, mock_requests_get, client, tmp_path, monkeypatch): def test_process_url_oserror_during_save(self, mock_stream, client, tmp_path, monkeypatch):
"""Test handling of OSError when saving file""" """Test handling of OSError when saving file"""
# Mock successful download # Mock successful download
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"}
mock_response.iter_content = Mock(return_value=[b"PDF"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
# Mock workdir to a non-existent path to trigger OSError # Mock workdir to a non-existent path to trigger OSError
from app.config import settings from app.config import settings
@@ -450,17 +499,24 @@ class TestURLUploadEndpoint:
# Restore original workdir # Restore original workdir
monkeypatch.setattr(settings, "workdir", original_workdir) monkeypatch.setattr(settings, "workdir", original_workdir)
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_unexpected_exception(self, mock_process_document, mock_requests_get, client): def test_process_url_unexpected_exception(self, mock_process_document, mock_stream, client):
"""Test handling of unexpected exceptions""" """Test handling of unexpected exceptions"""
# Mock successful download but process_document.delay raises unexpected error # Mock successful download but process_document.delay raises unexpected error
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"}
mock_response.iter_content = Mock(return_value=[b"PDF"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
# Mock process_document.delay to raise an unexpected exception # Mock process_document.delay to raise an unexpected exception
mock_process_document.delay.side_effect = RuntimeError("Unexpected processing error") mock_process_document.delay.side_effect = RuntimeError("Unexpected processing error")
@@ -471,17 +527,24 @@ class TestURLUploadEndpoint:
data = response.json() data = response.json()
assert "Unexpected error" in data["detail"] assert "Unexpected error" in data["detail"]
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_filename_without_extension(self, mock_process_document, mock_requests_get, client): def test_process_url_filename_without_extension(self, mock_process_document, mock_stream, client):
"""Test that files without extensions are handled correctly""" """Test that files without extensions are handled correctly"""
# Mock successful download # Mock successful download
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"}
mock_response.iter_content = Mock(return_value=[b"PDF"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
# Mock Celery task # Mock Celery task
mock_task = Mock() mock_task = Mock()
@@ -496,17 +559,24 @@ class TestURLUploadEndpoint:
# Should still work, just without extension # Should still work, just without extension
assert data["task_id"] == "test-task-id" assert data["task_id"] == "test-task-id"
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_empty_path_uses_download(self, mock_process_document, mock_requests_get, client): def test_process_url_empty_path_uses_download(self, mock_process_document, mock_stream, client):
"""Test that empty URL path defaults to 'download' filename""" """Test that empty URL path defaults to 'download' filename"""
# Mock successful download # Mock successful download
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"}
mock_response.iter_content = Mock(return_value=[b"PDF"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
# Mock Celery task # Mock Celery task
mock_task = Mock() mock_task = Mock()
@@ -560,17 +630,24 @@ class TestURLUploadEndpoint:
# Link-local address # Link-local address
assert is_private_ip("169.254.1.1") is True assert is_private_ip("169.254.1.1") is True
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_sanitizes_dangerous_filename(self, mock_process_document, mock_requests_get, client): def test_process_url_sanitizes_dangerous_filename(self, mock_process_document, mock_stream, client):
"""Test that dangerous filenames are sanitized""" """Test that dangerous filenames are sanitized"""
# Mock successful download # Mock successful download
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"}
mock_response.iter_content = Mock(return_value=[b"PDF"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
# Mock Celery task # Mock Celery task
mock_task = Mock() mock_task = Mock()
@@ -671,18 +748,25 @@ class TestURLUploadCoverageGaps:
assert validate_file_type("", "filename_without_extension") is False assert validate_file_type("", "filename_without_extension") is False
@patch("app.api.url_upload.sanitize_filename", return_value="") @patch("app.api.url_upload.sanitize_filename", return_value="")
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_sanitize_filename_returns_empty( def test_process_url_sanitize_filename_returns_empty(
self, mock_process_document, mock_requests_get, mock_sanitize, client self, mock_process_document, mock_stream, mock_sanitize, client
): ):
"""Test that when sanitize_filename returns empty string, filename defaults to 'download' (line 177)""" """Test that when sanitize_filename returns empty string, filename defaults to 'download' (line 177)"""
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"}
mock_response.iter_content = Mock(return_value=[b"PDF content"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF content"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
mock_task = Mock() mock_task = Mock()
mock_task.id = "test-task-id-sanitize" mock_task.id = "test-task-id-sanitize"
@@ -695,17 +779,26 @@ class TestURLUploadCoverageGaps:
# When sanitize_filename returns "", safe_filename defaults to "download" # When sanitize_filename returns "", safe_filename defaults to "download"
assert data["filename"] == "download" assert data["filename"] == "download"
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
@patch("app.api.url_upload.process_document") @patch("app.api.url_upload.process_document")
def test_process_url_skips_empty_chunks(self, mock_process_document, mock_requests_get, client): def test_process_url_skips_empty_chunks(self, mock_process_document, mock_stream, client):
"""Test that empty bytes chunks are skipped during download (line 234->233 branch)""" """Test that empty bytes chunks are skipped during download (line 234->233 branch)"""
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf"} mock_response.headers = {"Content-Type": "application/pdf"}
# Mix empty bytes (falsy) with real content - covers the `if chunk:` False branch # Mix empty bytes (falsy) with real content - covers the `if chunk:` False branch
mock_response.iter_content = Mock(return_value=[b"", b"PDF content", b""])
async def mock_aiter_bytes(chunk_size=None):
yield b""
yield b"PDF content"
yield b""
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
mock_task = Mock() mock_task = Mock()
mock_task.id = "test-task-id-chunks" mock_task.id = "test-task-id-chunks"
@@ -719,9 +812,9 @@ class TestURLUploadCoverageGaps:
@patch("app.api.url_upload.os.remove") @patch("app.api.url_upload.os.remove")
@patch("app.api.url_upload.os.path.exists", return_value=True) @patch("app.api.url_upload.os.path.exists", return_value=True)
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_oserror_cleanup_removes_existing_file( def test_process_url_oserror_cleanup_removes_existing_file(
self, mock_requests_get, mock_exists, mock_remove, client, tmp_path, monkeypatch self, mock_stream, mock_exists, mock_remove, client, tmp_path, monkeypatch
): ):
"""Test OSError handler removes the partial file when it exists (line 285)""" """Test OSError handler removes the partial file when it exists (line 285)"""
import os import os
@@ -735,12 +828,19 @@ class TestURLUploadCoverageGaps:
monkeypatch.setattr(settings, "workdir", str(non_existent)) monkeypatch.setattr(settings, "workdir", str(non_existent))
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"} mock_response.headers = {"Content-Type": "application/pdf", "Content-Length": "100"}
mock_response.iter_content = Mock(return_value=[b"PDF"])
async def mock_aiter_bytes(chunk_size=None):
yield b"PDF"
mock_response.aiter_bytes = mock_aiter_bytes
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
response = client.post("/api/process-url", json={"url": "https://example.com/file.pdf"}) response = client.post("/api/process-url", json={"url": "https://example.com/file.pdf"})
@@ -750,14 +850,17 @@ class TestURLUploadCoverageGaps:
mock_remove.assert_called_once() mock_remove.assert_called_once()
@patch("app.api.url_upload.validate_file_type", side_effect=ValueError("unexpected internal error")) @patch("app.api.url_upload.validate_file_type", side_effect=ValueError("unexpected internal error"))
@patch("app.api.url_upload.requests.get") @patch("app.api.url_upload.httpx.AsyncClient.stream")
def test_process_url_unexpected_exception_with_no_file_created(self, mock_requests_get, mock_validate, client): def test_process_url_unexpected_exception_with_no_file_created(self, mock_stream, mock_validate, client):
"""Test unexpected exception before target_path is assigned; no file cleanup attempted (line 291->293)""" """Test unexpected exception before target_path is assigned; no file cleanup attempted (line 291->293)"""
mock_response = Mock() mock_response = AsyncMock()
mock_response.status_code = 200 mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/pdf"} mock_response.headers = {"Content-Type": "application/pdf"}
mock_response.raise_for_status = Mock() mock_response.raise_for_status = Mock()
mock_requests_get.return_value = mock_response
mock_context = AsyncMock()
mock_context.__aenter__.return_value = mock_response
mock_stream.return_value = mock_context
response = client.post("/api/process-url", json={"url": "https://example.com/file.pdf"}) response = client.post("/api/process-url", json={"url": "https://example.com/file.pdf"})