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:
+56
-56
@@ -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)}")
|
||||||
|
|
||||||
|
|||||||
@@ -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())
|
||||||
@@ -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
@@ -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"})
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user