Merge pull request #562 from christianlouis/copilot/refactor-dynamic-routing-user-destinations

fix(tasks): register upload_to_user_integration in Celery and add handler coverage
This commit is contained in:
Christian Krakau-Louis
2026-03-09 01:05:36 +01:00
committed by GitHub
8 changed files with 2549 additions and 60 deletions
+1
View File
@@ -38,6 +38,7 @@ from app.tasks.upload_to_onedrive import upload_to_onedrive # noqa: F401
from app.tasks.upload_to_paperless import upload_to_paperless # noqa: F401
from app.tasks.upload_to_s3 import upload_to_s3 # noqa: F401
from app.tasks.upload_to_sftp import upload_to_sftp # noqa: F401
from app.tasks.upload_to_user_integration import upload_to_user_integration # noqa: F401
from app.tasks.upload_to_webdav import upload_to_webdav # noqa: F401
from app.tasks.upload_with_rclone import send_to_all_rclone_destinations, upload_with_rclone # noqa: F401
from app.tasks.uptime_kuma_tasks import ping_uptime_kuma # noqa: F401
+50 -17
View File
@@ -10,8 +10,13 @@ from app.database import SessionLocal
from app.models import FileRecord
from app.tasks.retry_config import BaseTaskWithRetry
# Import the aggregator task and validator
from app.tasks.send_to_all import get_configured_services_from_validator, send_to_all_destinations
# Import the aggregator tasks and validator
from app.tasks.send_to_all import (
get_configured_services_from_validator,
get_user_destination_count,
send_to_all_destinations,
send_to_user_destinations,
)
# Import database and logging utils from main
from app.utils import log_task_progress
@@ -26,8 +31,16 @@ logger = logging.getLogger(__name__)
def finalize_document_storage(self, original_file: str, processed_file: str, metadata: dict, file_id: int = None):
"""
Final storage step after embedding metadata.
We will now call 'send_to_all_destinations' to push the final PDF to Dropbox/Nextcloud/Paperless.
After uploading, send a notification about the processed file.
Routes the processed document to the appropriate destination(s):
1. If the document has an identified owner and that owner has active
DESTINATION UserIntegrations, the file is uploaded to each of those
integrations (user-specific routing).
2. Otherwise the file is forwarded to the globally-configured destinations
via :func:`send_to_all_destinations` (system-wide fallback).
After queuing uploads, optional PDF/A archival conversion and embedding
computation are triggered, and a completion notification is sent.
"""
task_id = self.request.id
logger.info(f"[{task_id}] Finalizing document storage for {processed_file}")
@@ -41,7 +54,8 @@ def finalize_document_storage(self, original_file: str, processed_file: str, met
file_id=file_id,
)
# Get file_id from database if not provided (fallback logic)
# 2. Resolve file_id and owner_id from the database
owner_id = None
if file_id is None:
with SessionLocal() as db:
# Only as a last resort, try to find by exact match on local_filename
@@ -49,32 +63,52 @@ def finalize_document_storage(self, original_file: str, processed_file: str, met
file_record = db.query(FileRecord).filter(FileRecord.local_filename == tmp_path).first()
if file_record:
file_id = file_record.id
owner_id = file_record.owner_id
else:
with SessionLocal() as db:
file_record = db.query(FileRecord).filter(FileRecord.id == file_id).first()
if file_record:
owner_id = file_record.owner_id
# 2. Determine Configured Destinations
# This is needed for the notification message later
# 3. Determine configured destinations for notification
configured_destinations = []
try:
configured_services = get_configured_services_from_validator()
# Get list of service names that are configured
for service_name, is_configured in configured_services.items():
if is_configured:
# Format service names for display
display_name = service_name.replace("_", " ").title()
configured_destinations.append(display_name)
except Exception as e:
logger.warning(f"[WARNING] Could not determine configured destinations: {e}")
configured_destinations = ["configured destinations"]
# 3. Queue Uploads
logger.info(f"[{task_id}] Queueing uploads to all destinations")
# 4. Queue Uploads — prefer user-specific destinations when available
log_task_progress(
task_id, "finalize_document_storage", "success", "Queuing uploads to destinations", file_id=file_id
)
# Note: send_to_all_destinations is asynchronous and queues upload tasks
send_to_all_destinations.delay(processed_file, True, file_id)
user_dest_count = 0
if owner_id:
try:
user_dest_count = get_user_destination_count(owner_id)
except Exception as e:
logger.warning("[%s] Could not query user destination count for owner=%s: %s", task_id, owner_id, e)
# 3a. Trigger PDF/A archival conversion if enabled (from feature branch)
if owner_id and user_dest_count > 0:
# User has configured their own destinations → use those exclusively
logger.info(
"[%s] Routing to %d user-specific destination(s) for owner=%s",
task_id,
user_dest_count,
owner_id,
)
send_to_user_destinations.delay(processed_file, owner_id, file_id)
else:
# No user-specific destinations → fall back to global configuration
logger.info("[%s] No user-specific destinations found; using global destinations", task_id)
send_to_all_destinations.delay(processed_file, True, file_id)
# 4a. Trigger PDF/A archival conversion if enabled
if settings.enable_pdfa_conversion:
try:
from app.tasks.convert_to_pdfa import convert_to_pdfa
@@ -84,7 +118,7 @@ def finalize_document_storage(self, original_file: str, processed_file: str, met
except Exception as e:
logger.warning(f"[{task_id}] Could not queue PDF/A conversion: {e}")
# 3b. Queue embedding computation (from main branch)
# 4b. Queue embedding computation
if file_id is not None:
try:
from app.tasks.compute_embedding import compute_document_embedding
@@ -94,9 +128,8 @@ def finalize_document_storage(self, original_file: str, processed_file: str, met
except Exception as e:
logger.warning(f"[{task_id}] Could not queue embedding task: {e}")
# 4. Send Notification
# 5. Send Notification
try:
# Get file information
file_size = os.path.getsize(processed_file) if os.path.exists(processed_file) else 0
filename = os.path.basename(processed_file)
+121 -1
View File
@@ -6,7 +6,7 @@ import os
from app.celery_app import celery
from app.config import settings
from app.database import SessionLocal
from app.models import FileRecord
from app.models import FileRecord, IntegrationDirection, UserIntegration
from app.tasks.retry_config import BaseTaskWithRetry
from app.tasks.upload_to_dropbox import upload_to_dropbox
from app.tasks.upload_to_email import upload_to_email
@@ -261,3 +261,123 @@ def send_to_all_destinations(self, file_path: str, use_validator=True, file_id:
log_task_progress(task_id, "send_to_all_destinations", "success", f"Queued {queued_count} uploads", file_id=file_id)
return {"status": "Queued", "file_path": file_path, "tasks": results}
@celery.task(base=BaseTaskWithRetry, bind=True)
def send_to_user_destinations(self, file_path: str, owner_id: str, file_id: int | None = None):
"""Dispatch uploads to all active DESTINATION UserIntegrations for *owner_id*.
This is the user-specific counterpart of :func:`send_to_all_destinations`.
It queries the ``user_integrations`` table for records where:
* ``owner_id`` matches the document owner,
* ``direction == "DESTINATION"``, and
* ``is_active == True``.
One :func:`upload_to_user_integration` Celery task is queued for each
matching integration so that uploads proceed asynchronously and
independently.
Args:
file_path: Absolute path to the processed document file.
owner_id: The stable user identifier from ``FileRecord.owner_id``.
file_id: Optional ``FileRecord.id`` used for progress logging.
Returns:
A dict summarising how many integrations were queued.
"""
from app.tasks.upload_to_user_integration import upload_to_user_integration
task_id = self.request.id
filename = os.path.basename(file_path)
if not os.path.exists(file_path):
error_msg = f"File not found: {file_path}"
logger.error("[%s] %s", task_id, error_msg)
log_task_progress(task_id, "send_to_user_destinations", "failure", error_msg, file_id=file_id)
raise FileNotFoundError(error_msg)
logger.info("[%s] Sending %s to user destinations for owner=%s", task_id, filename, owner_id)
log_task_progress(
task_id,
"send_to_user_destinations",
"in_progress",
f"Distributing {filename} to user integrations",
file_id=file_id,
)
with SessionLocal() as db:
integrations = (
db.query(UserIntegration)
.filter(
UserIntegration.owner_id == owner_id,
UserIntegration.direction == IntegrationDirection.DESTINATION,
UserIntegration.is_active.is_(True),
)
.all()
)
# Snapshot the IDs so we don't keep the session open
integration_ids = [(i.id, i.name, i.integration_type) for i in integrations]
queued = 0
task_results: dict[str, str] = {}
for int_id, int_name, int_type in integration_ids:
logger.info("[%s] Queueing upload for integration %d (%s '%s')", task_id, int_id, int_type, int_name)
log_task_progress(
task_id,
f"queue_user_integration_{int_id}",
"in_progress",
f"Queueing upload to {int_type} '{int_name}'",
file_id=file_id,
)
try:
celery_task = upload_to_user_integration.delay(file_path, int_id, file_id)
task_results[f"integration_{int_id}_task_id"] = celery_task.id
queued += 1
log_task_progress(
task_id,
f"queue_user_integration_{int_id}",
"success",
f"Queued upload to {int_type} '{int_name}'",
file_id=file_id,
)
except Exception as exc: # noqa: BLE001
error_msg = str(exc)
logger.error("[%s] Failed to queue upload for integration %d: %s", task_id, int_id, error_msg)
task_results[f"integration_{int_id}_error"] = error_msg
log_task_progress(
task_id,
f"queue_user_integration_{int_id}",
"failure",
f"Failed to queue {int_type} '{int_name}': {error_msg}",
file_id=file_id,
)
logger.info("[%s] Queued %d user-integration upload(s) for owner=%s", task_id, queued, owner_id)
log_task_progress(
task_id,
"send_to_user_destinations",
"success",
f"Queued {queued} user-integration upload(s)",
file_id=file_id,
)
return {"status": "Queued", "file_path": file_path, "queued": queued, "tasks": task_results}
def get_user_destination_count(owner_id: str) -> int:
"""Return the number of active DESTINATION integrations for *owner_id*.
A count of zero means no user-specific destinations are configured and
the caller should fall back to the global :func:`send_to_all_destinations`.
"""
with SessionLocal() as db:
return (
db.query(UserIntegration)
.filter(
UserIntegration.owner_id == owner_id,
UserIntegration.direction == IntegrationDirection.DESTINATION,
UserIntegration.is_active.is_(True),
)
.count()
)
+731
View File
@@ -0,0 +1,731 @@
#!/usr/bin/env python3
"""
Upload dispatcher for user-specific destination integrations.
This module provides a Celery task that uploads a processed document to a
specific ``UserIntegration`` record using that integration's own stored
config and decrypted credentials — instead of the global application settings.
It is the per-destination counterpart of :func:`send_to_all_destinations`
and is dispatched by :func:`send_to_user_destinations` once per active
DESTINATION integration that belongs to the document's owner.
Credential shapes per integration type (mirrors the UserIntegration docstring):
DROPBOX credentials = {"refresh_token", "app_key", "app_secret"}
config = {"folder": "/DocuElevate"}
S3 credentials = {"access_key_id", "secret_access_key"}
config = {"bucket", "region", "endpoint_url", "folder_prefix"}
GOOGLE_DRIVE
OAuth credentials = {"client_id", "client_secret", "refresh_token"}
config = {"folder_id"}
SA credentials = {"credentials_json"}
config = {"folder_id"}
ONEDRIVE credentials = {"client_id", "client_secret", "refresh_token"}
config = {"folder_path", "tenant_id"}
WEBDAV /
NEXTCLOUD credentials = {"username", "password"}
config = {"url", "folder"}
FTP credentials = {"password"}
config = {"host", "username", "port", "folder", "use_tls"}
SFTP credentials = {"password"} or {"private_key"}
config = {"host", "username", "port", "folder"}
EMAIL credentials = {"password"}
config = {"host", "username", "port", "recipient",
"use_tls", "sender_name"}
PAPERLESS credentials = {"api_token"}
config = {"host"}
RCLONE credentials = {"rclone_conf"} (full rclone config file text)
config = {"remote": "myremote:", "folder": "dest/path"}
"""
import ftplib # nosec B402
import json
import logging
import os
import subprocess # nosec B404
from datetime import datetime, timezone
from typing import Any
from urllib.parse import urljoin
from app.celery_app import celery
from app.database import SessionLocal
from app.models import IntegrationType, UserIntegration
from app.tasks.retry_config import UploadTaskWithRetry
from app.utils.encryption import decrypt_value
from app.utils.logging import log_task_progress
logger = logging.getLogger(__name__)
# Maximum characters to store in UserIntegration.last_error
_MAX_ERROR_LENGTH = 500
# ---------------------------------------------------------------------------
# Per-type upload helpers
# ---------------------------------------------------------------------------
def _upload_dropbox(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Upload *file_path* to Dropbox using per-user OAuth credentials."""
import dropbox
app_key = creds.get("app_key") or ""
app_secret = creds.get("app_secret") or ""
refresh_token = creds.get("refresh_token") or ""
if not (app_key and app_secret and refresh_token):
raise ValueError("Dropbox integration is missing app_key, app_secret or refresh_token in credentials")
dbx = dropbox.Dropbox(app_key=app_key, app_secret=app_secret, oauth2_refresh_token=refresh_token)
remote_folder = cfg.get("folder", "/DocuElevate").rstrip("/")
filename = os.path.basename(file_path)
remote_path = f"{remote_folder}/{filename}"
if not remote_path.startswith("/"):
remote_path = "/" + remote_path
file_size = os.path.getsize(file_path)
with open(file_path, "rb") as fh:
if file_size > 10 * 1024 * 1024:
chunk_size = 4 * 1024 * 1024
session_start = dbx.files_upload_session_start(fh.read(chunk_size))
cursor = dropbox.files.UploadSessionCursor(session_start.session_id, fh.tell())
while fh.tell() < file_size:
if (file_size - fh.tell()) <= chunk_size:
dbx.files_upload_session_finish(
fh.read(chunk_size),
cursor,
dropbox.files.CommitInfo(path=remote_path, mode=dropbox.files.WriteMode.overwrite),
)
else:
dbx.files_upload_session_append_v2(fh.read(chunk_size), cursor)
cursor.offset = fh.tell()
else:
dbx.files_upload(fh.read(), remote_path, mode=dropbox.files.WriteMode.overwrite)
logger.info("[%s] Dropbox upload complete: %s", task_id, remote_path)
return {"status": "Completed", "dropbox_path": remote_path}
def _upload_s3(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Upload *file_path* to Amazon S3 (or S3-compatible) using per-user credentials."""
import boto3
from botocore.exceptions import ClientError
bucket = cfg.get("bucket") or ""
region = cfg.get("region") or "us-east-1"
endpoint_url = cfg.get("endpoint_url") or None
folder_prefix = cfg.get("folder_prefix") or ""
storage_class = cfg.get("storage_class") or "STANDARD"
access_key = creds.get("access_key_id") or ""
secret_key = creds.get("secret_access_key") or ""
if not bucket:
raise ValueError("S3 integration is missing bucket in config")
if not (access_key and secret_key):
raise ValueError("S3 integration is missing access_key_id or secret_access_key in credentials")
client_kwargs: dict[str, Any] = {
"region_name": region,
"aws_access_key_id": access_key,
"aws_secret_access_key": secret_key,
}
if endpoint_url:
client_kwargs["endpoint_url"] = endpoint_url
s3 = boto3.client("s3", **client_kwargs)
filename = os.path.basename(file_path)
s3_key = f"{folder_prefix.rstrip('/')}/{filename}" if folder_prefix else filename
try:
s3.upload_file(file_path, bucket, s3_key, ExtraArgs={"StorageClass": storage_class})
except ClientError as exc:
raise RuntimeError(f"S3 upload failed: {exc}") from exc
logger.info("[%s] S3 upload complete: s3://%s/%s", task_id, bucket, s3_key)
return {"status": "Completed", "s3_bucket": bucket, "s3_key": s3_key}
def _upload_google_drive(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Upload *file_path* to Google Drive using per-user OAuth or service-account credentials."""
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload
folder_id = cfg.get("folder_id") or ""
filename = os.path.basename(file_path)
# Prefer OAuth credentials (client_id + client_secret + refresh_token)
client_id = creds.get("client_id") or ""
client_secret = creds.get("client_secret") or ""
refresh_token = creds.get("refresh_token") or ""
credentials_json = creds.get("credentials_json") or ""
if client_id and client_secret and refresh_token:
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials as OAuthCredentials
google_creds = OAuthCredentials(
None,
refresh_token=refresh_token,
token_uri="https://oauth2.googleapis.com/token",
client_id=client_id,
client_secret=client_secret,
scopes=["https://www.googleapis.com/auth/drive.file"],
)
google_creds.refresh(Request())
service = build("drive", "v3", credentials=google_creds)
elif credentials_json:
from google.oauth2.service_account import Credentials as SACredentials
creds_dict = json.loads(credentials_json)
sa_creds = SACredentials.from_service_account_info(creds_dict, scopes=["https://www.googleapis.com/auth/drive"])
service = build("drive", "v3", credentials=sa_creds)
else:
raise ValueError("Google Drive integration requires either OAuth credentials or credentials_json")
file_metadata: dict[str, Any] = {"name": filename}
if folder_id:
file_metadata["parents"] = [folder_id]
media = MediaFileUpload(file_path, mimetype="application/pdf", resumable=True)
file_obj = service.files().create(body=file_metadata, media_body=media, fields="id,name,webViewLink").execute()
gdrive_id = file_obj.get("id")
web_link = file_obj.get("webViewLink")
logger.info("[%s] Google Drive upload complete: %s (%s)", task_id, gdrive_id, web_link)
return {"status": "Completed", "google_drive_file_id": gdrive_id, "web_link": web_link}
def _upload_onedrive(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Upload *file_path* to OneDrive using per-user MSAL credentials."""
import urllib.parse
import msal
import requests as _requests
client_id = creds.get("client_id") or ""
client_secret = creds.get("client_secret") or ""
refresh_token = creds.get("refresh_token") or ""
tenant = cfg.get("tenant_id") or "common"
folder_path = cfg.get("folder_path") or ""
if not (client_id and client_secret):
raise ValueError("OneDrive integration is missing client_id or client_secret in credentials")
scopes = ["https://graph.microsoft.com/.default"]
msal_app = msal.ConfidentialClientApplication(
client_id=client_id,
client_credential=client_secret,
authority=f"https://login.microsoftonline.com/{tenant}",
)
if refresh_token:
token_resp = msal_app.acquire_token_by_refresh_token(refresh_token=refresh_token, scopes=scopes)
else:
token_resp = msal_app.acquire_token_for_client(scopes=scopes)
if "access_token" not in token_resp:
raise ValueError(f"OneDrive token acquisition failed: {token_resp.get('error_description', 'unknown')}")
access_token = token_resp["access_token"]
filename = os.path.basename(file_path)
# Build upload-session URL
if folder_path:
folder_path = folder_path.strip("/")
encoded_path = "/".join(urllib.parse.quote(p) for p in folder_path.split("/"))
encoded_file = urllib.parse.quote(filename)
item_path = f"/root:/{encoded_path}/{encoded_file}:/createUploadSession"
else:
encoded_file = urllib.parse.quote(filename)
item_path = f"/root:/{encoded_file}:/createUploadSession"
session_url = f"https://graph.microsoft.com/v1.0/me/drive{item_path}"
headers = {"Authorization": f"Bearer {access_token}", "Content-Type": "application/json"}
resp = _requests.post(
session_url, headers=headers, json={"item": {"@microsoft.graph.conflictBehavior": "replace"}}, timeout=30
)
resp.raise_for_status()
upload_url = resp.json()["uploadUrl"]
file_size = os.path.getsize(file_path)
chunk_size = 10 * 1024 * 1024
with open(file_path, "rb") as fh:
chunk_num = 0
while True:
chunk = fh.read(chunk_size)
if not chunk:
break
start = chunk_num * chunk_size
end = start + len(chunk) - 1
upload_headers = {
"Content-Length": str(len(chunk)),
"Content-Range": f"bytes {start}-{end}/{file_size}",
}
upload_resp = _requests.put(upload_url, headers=upload_headers, data=chunk, timeout=120)
if upload_resp.status_code not in (201, 202):
raise RuntimeError(f"OneDrive chunk upload failed: {upload_resp.status_code}")
chunk_num += 1
logger.info("[%s] OneDrive upload complete: %s/%s", task_id, folder_path, filename)
return {"status": "Completed", "onedrive_folder": folder_path, "filename": filename}
def _upload_webdav(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Upload *file_path* to a WebDAV server using per-user credentials."""
import requests as _requests
url = cfg.get("url") or ""
folder = cfg.get("folder") or ""
username = creds.get("username") or ""
password = creds.get("password") or ""
verify_ssl = cfg.get("verify_ssl", True)
if not url:
raise ValueError("WebDAV integration is missing url in config")
filename = os.path.basename(file_path)
folder = folder.lstrip("/")
target = urljoin(url.rstrip("/") + "/", folder)
if not target.endswith("/"):
target += "/"
dest = urljoin(target, filename)
with open(file_path, "rb") as fh:
resp = _requests.put(dest, auth=(username, password), data=fh, verify=verify_ssl, timeout=120)
if resp.status_code not in (200, 201, 204):
raise RuntimeError(f"WebDAV upload failed: {resp.status_code} {resp.text[:200]}")
logger.info("[%s] WebDAV upload complete: %s", task_id, dest)
return {"status": "Completed", "webdav_url": dest}
def _upload_nextcloud(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Upload *file_path* to Nextcloud (WebDAV) using per-user credentials."""
# Nextcloud uses WebDAV under the hood; reuse the WebDAV helper.
return _upload_webdav(file_path, cfg, creds, task_id)
def _upload_ftp(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Upload *file_path* to an FTP/FTPS server using per-user credentials."""
host = cfg.get("host") or ""
port = int(cfg.get("port") or 21)
username = cfg.get("username") or ""
folder = cfg.get("folder") or ""
use_tls = cfg.get("use_tls", True)
password = creds.get("password") or ""
filename = os.path.basename(file_path)
if not host:
raise ValueError("FTP integration is missing host in config")
ftp: ftplib.FTP
if use_tls:
ftp = ftplib.FTP_TLS() # nosec B321 # noqa: S321
ftp.connect(host=host, port=port)
ftp.login(user=username, passwd=password)
ftp.prot_p()
else:
ftp = ftplib.FTP() # nosec B321 # noqa: S321
ftp.connect(host=host, port=port)
ftp.login(user=username, passwd=password)
if folder:
folder_stripped = folder.lstrip("/")
try:
ftp.cwd(folder_stripped)
except ftplib.error_perm:
parts = folder_stripped.split("/")
current = ""
for part in parts:
if not part:
continue
current += f"/{part}"
try:
ftp.cwd(current)
except ftplib.error_perm:
ftp.mkd(current)
ftp.cwd(current)
with open(file_path, "rb") as fh:
ftp.storbinary(f"STOR {filename}", fh)
ftp.quit()
logger.info("[%s] FTP upload complete: %s/%s", task_id, host, filename)
return {"status": "Completed", "ftp_host": host, "filename": filename}
def _upload_sftp(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Upload *file_path* to an SFTP server using per-user credentials."""
import paramiko
host = cfg.get("host") or ""
port = int(cfg.get("port") or 22)
username = cfg.get("username") or ""
folder = cfg.get("folder") or ""
password = creds.get("password") or ""
private_key_text = creds.get("private_key") or ""
filename = os.path.basename(file_path)
if not host:
raise ValueError("SFTP integration is missing host in config")
ssh = paramiko.SSHClient()
ssh.load_system_host_keys()
ssh.set_missing_host_key_policy(paramiko.RejectPolicy())
connect_kwargs: dict[str, Any] = {"hostname": host, "port": port, "username": username}
if private_key_text:
import io
pkey = paramiko.RSAKey.from_private_key(io.StringIO(private_key_text))
connect_kwargs["pkey"] = pkey
elif password:
connect_kwargs["password"] = password
else:
raise ValueError("SFTP integration requires password or private_key in credentials")
ssh.connect(**connect_kwargs)
sftp = ssh.open_sftp()
remote_path = f"{folder.rstrip('/')}/{filename}" if folder else filename
if folder and folder.startswith("/") and not remote_path.startswith("/"):
remote_path = "/" + remote_path
sftp.put(file_path, remote_path)
sftp.close()
ssh.close()
logger.info("[%s] SFTP upload complete: %s:%s", task_id, host, remote_path)
return {"status": "Completed", "sftp_host": host, "sftp_path": remote_path}
def _upload_paperless(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Upload *file_path* to a Paperless-ngx instance using per-user API token."""
import time
import requests as _requests
host = (cfg.get("host") or "").rstrip("/")
api_token = creds.get("api_token") or ""
filename = os.path.basename(file_path)
if not host:
raise ValueError("Paperless integration is missing host in config")
if not api_token:
raise ValueError("Paperless integration is missing api_token in credentials")
headers = {"Authorization": f"Token {api_token}"}
post_url = f"{host}/api/documents/post_document/"
with open(file_path, "rb") as fh:
resp = _requests.post(
post_url,
headers=headers,
files={"document": (filename, fh, "application/pdf")},
data={"title": filename},
timeout=120,
)
resp.raise_for_status()
raw_task_id = resp.text.strip().strip('"').strip("'")
# Poll for completion (up to 30 s)
task_url = f"{host}/api/tasks/"
doc_id = None
for _ in range(10):
time.sleep(3)
try:
poll_resp = _requests.get(task_url, headers=headers, params={"task_id": raw_task_id}, timeout=30)
poll_resp.raise_for_status()
tasks_data = poll_resp.json()
if isinstance(tasks_data, dict) and "results" in tasks_data:
tasks_data = tasks_data["results"]
if tasks_data:
info = tasks_data[0]
status = info.get("status")
if status == "SUCCESS":
doc_id = info.get("related_document")
break
elif status == "FAILURE":
raise RuntimeError(f"Paperless processing failed: {info.get('result')}")
except RuntimeError:
raise
except Exception as poll_exc:
logger.warning("[%s] Paperless poll error: %s", task_id, poll_exc)
logger.info("[%s] Paperless upload complete: doc_id=%s", task_id, doc_id)
return {"status": "Completed", "paperless_host": host, "paperless_document_id": doc_id}
def _upload_email(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Send *file_path* as an email attachment using per-user SMTP credentials."""
import smtplib
from email.mime.application import MIMEApplication
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
host = cfg.get("host") or ""
port = int(cfg.get("port") or 587)
username = cfg.get("username") or ""
recipient = cfg.get("recipient") or ""
use_tls = cfg.get("use_tls", True)
sender_name = cfg.get("sender_name") or "DocuElevate"
password = creds.get("password") or ""
filename = os.path.basename(file_path)
if not (host and recipient):
raise ValueError("Email integration is missing host or recipient in config")
msg = MIMEMultipart()
msg["From"] = f"{sender_name} <{username}>" if username else sender_name
msg["To"] = recipient
msg["Subject"] = f"Document: {filename}"
msg.attach(MIMEText(f"Please find the attached document: {filename}", "plain"))
with open(file_path, "rb") as fh:
part = MIMEApplication(fh.read(), Name=filename)
part["Content-Disposition"] = f'attachment; filename="{filename}"'
msg.attach(part)
if use_tls:
import ssl
tls_context = ssl.create_default_context()
with smtplib.SMTP(host, port, timeout=30) as smtp:
smtp.starttls(context=tls_context)
if username and password:
smtp.login(username, password)
smtp.sendmail(msg["From"], [recipient], msg.as_string())
else:
# Plaintext SMTP — only use when explicitly configured and TLS is unavailable.
# Credentials and content will be transmitted without encryption.
with smtplib.SMTP(host, port, timeout=30) as smtp: # nosec B608
if username and password:
smtp.login(username, password)
smtp.sendmail(msg["From"], [recipient], msg.as_string())
logger.info("[%s] Email upload complete: sent to %s", task_id, recipient)
return {"status": "Completed", "recipient": recipient}
def _upload_rclone(file_path: str, cfg: dict[str, Any], creds: dict[str, Any], task_id: str) -> dict[str, Any]:
"""Copy *file_path* to an rclone remote using per-user rclone config."""
import re
import tempfile
remote = cfg.get("remote") or ""
folder = cfg.get("folder") or ""
rclone_conf_text = creds.get("rclone_conf") or ""
filename = os.path.basename(file_path)
if not remote:
raise ValueError("Rclone integration is missing remote in config")
if not rclone_conf_text:
raise ValueError("Rclone integration is missing rclone_conf in credentials")
# Validate remote and folder to prevent shell metacharacter injection.
# rclone remote names are alphanumeric + hyphens/underscores followed by ':'.
# folder paths must not contain shell-dangerous characters.
_SAFE_REMOTE_RE = re.compile(r"^[A-Za-z0-9_\-]+:(/[A-Za-z0-9_.@\-/ ]*)?$")
_SAFE_FOLDER_RE = re.compile(r"^[A-Za-z0-9_.@\-/ ]*$")
if not _SAFE_REMOTE_RE.match(remote):
raise ValueError(f"Rclone remote contains unsafe characters: {remote!r}")
if folder and not _SAFE_FOLDER_RE.match(folder):
raise ValueError(f"Rclone folder contains unsafe characters: {folder!r}")
# Write the user's rclone config to a temp file so we don't touch the system config
with tempfile.NamedTemporaryFile(mode="w", suffix=".conf", delete=False) as tmp_conf:
tmp_conf.write(rclone_conf_text)
conf_path = tmp_conf.name
dest = f"{remote.rstrip('/')}/{folder.strip('/')}/{filename}" if folder else f"{remote.rstrip('/')}/{filename}"
dest = dest.replace("//", "/")
try:
result = subprocess.run( # nosec B603 # noqa: S603 S607
["rclone", "copyto", f"--config={conf_path}", file_path, dest], # noqa: S603 S607
capture_output=True,
text=True,
timeout=300,
check=False,
)
if result.returncode != 0:
raise RuntimeError(f"rclone exited {result.returncode}: {result.stderr[:300]}")
finally:
os.unlink(conf_path)
logger.info("[%s] Rclone upload complete: %s", task_id, dest)
return {"status": "Completed", "rclone_dest": dest}
# Map IntegrationType → upload helper
_UPLOAD_HANDLERS = {
IntegrationType.DROPBOX: _upload_dropbox,
IntegrationType.S3: _upload_s3,
IntegrationType.GOOGLE_DRIVE: _upload_google_drive,
IntegrationType.ONEDRIVE: _upload_onedrive,
IntegrationType.WEBDAV: _upload_webdav,
IntegrationType.NEXTCLOUD: _upload_nextcloud,
IntegrationType.FTP: _upload_ftp,
IntegrationType.SFTP: _upload_sftp,
IntegrationType.PAPERLESS: _upload_paperless,
IntegrationType.EMAIL: _upload_email,
IntegrationType.RCLONE: _upload_rclone,
}
# ---------------------------------------------------------------------------
# Celery task
# ---------------------------------------------------------------------------
@celery.task(base=UploadTaskWithRetry, bind=True)
def upload_to_user_integration(self, file_path: str, integration_id: int, file_id: int | None = None) -> dict[str, Any]:
"""Upload *file_path* to the destination described by the given UserIntegration record.
This task is dispatched once per active DESTINATION UserIntegration that
belongs to a document's owner. Credentials are decrypted at runtime so
they never travel across the Celery message bus in plaintext.
Args:
file_path: Absolute path to the processed document file.
integration_id: Primary key of the ``UserIntegration`` record.
file_id: Optional ``FileRecord.id`` used for progress logging.
Returns:
A dict with at least ``{"status": "Completed", ...}`` on success.
Raises:
FileNotFoundError: When *file_path* does not exist.
ValueError: When the integration record is not found or has missing config.
RuntimeError: When the underlying upload operation fails.
"""
task_id = self.request.id
filename = os.path.basename(file_path)
log_task_progress(
task_id,
f"upload_to_user_integration_{integration_id}",
"in_progress",
f"Uploading {filename} to integration {integration_id}",
file_id=file_id,
)
if not os.path.exists(file_path):
error_msg = f"File not found: {file_path}"
logger.error("[%s] %s", task_id, error_msg)
log_task_progress(
task_id, f"upload_to_user_integration_{integration_id}", "failure", error_msg, file_id=file_id
)
raise FileNotFoundError(error_msg)
with SessionLocal() as db:
integration: UserIntegration | None = (
db.query(UserIntegration).filter(UserIntegration.id == integration_id).first()
)
if integration is None:
error_msg = f"UserIntegration {integration_id} not found"
logger.error("[%s] %s", task_id, error_msg)
log_task_progress(
task_id, f"upload_to_user_integration_{integration_id}", "failure", error_msg, file_id=file_id
)
raise ValueError(error_msg)
itype = integration.integration_type
int_name = integration.name
owner_id = integration.owner_id
# Parse config (non-sensitive) and decrypt credentials (sensitive)
try:
cfg: dict[str, Any] = json.loads(integration.config) if integration.config else {}
except json.JSONDecodeError as exc:
raise ValueError(f"Integration {integration_id} has invalid JSON in config: {exc}") from exc
try:
raw_creds = decrypt_value(integration.credentials) if integration.credentials else None
creds: dict[str, Any] = json.loads(raw_creds) if raw_creds else {}
except json.JSONDecodeError as exc:
raise ValueError(f"Integration {integration_id} has invalid JSON in credentials: {exc}") from exc
handler = _UPLOAD_HANDLERS.get(itype)
if handler is None:
error_msg = f"No upload handler registered for integration type '{itype}' (integration {integration_id})"
logger.warning("[%s] %s", task_id, error_msg)
log_task_progress(
task_id, f"upload_to_user_integration_{integration_id}", "skipped", error_msg, file_id=file_id
)
return {"status": "Skipped", "reason": error_msg}
logger.info(
"[%s] Uploading %s via %s integration '%s' (id=%d, owner=%s)",
task_id,
filename,
itype,
int_name,
integration_id,
owner_id,
)
try:
result = handler(file_path, cfg, creds, task_id)
# Update last_used_at on success
with SessionLocal() as db:
integ = db.query(UserIntegration).filter(UserIntegration.id == integration_id).first()
if integ:
integ.last_used_at = datetime.now(timezone.utc)
integ.last_error = None
db.commit()
log_task_progress(
task_id,
f"upload_to_user_integration_{integration_id}",
"success",
f"Uploaded to {itype} '{int_name}': {filename}",
file_id=file_id,
)
return result
except Exception as exc:
error_msg = str(exc)[:_MAX_ERROR_LENGTH]
logger.error(
"[%s] Upload to integration %d (%s '%s') failed: %s",
task_id,
integration_id,
itype,
int_name,
error_msg,
)
# Persist error for operator visibility
try:
with SessionLocal() as db:
integ = db.query(UserIntegration).filter(UserIntegration.id == integration_id).first()
if integ:
integ.last_used_at = datetime.now(timezone.utc)
integ.last_error = error_msg
db.commit()
except Exception as db_exc: # noqa: BLE001
logger.warning("[%s] Could not persist last_error for integration %d: %s", task_id, integration_id, db_exc)
log_task_progress(
task_id,
f"upload_to_user_integration_{integration_id}",
"failure",
f"Upload to {itype} '{int_name}' failed: {error_msg}",
file_id=file_id,
)
raise
+72
View File
@@ -295,6 +295,78 @@ with SessionLocal() as db:
pass
```
## User-Specific Destination Routing
### Overview
When a document has an identified owner (non-anonymous user), DocuElevate
routes the processed file to **that user's own configured destinations** instead
of the system-wide global destinations. This enables true multi-tenant
operation: each user's documents are stored where *they* configured, using
*their* OAuth tokens or API credentials.
### Routing Decision
The routing decision is made in `finalize_document_storage` after all
processing steps are complete:
```
Document owner has active DESTINATION integrations?
├── YES → send_to_user_destinations (user-specific routing)
└── NO → send_to_all_destinations (global fallback)
```
"Active DESTINATION integrations" means rows in the `user_integrations` table
where `owner_id` matches, `direction = "DESTINATION"`, and `is_active = True`.
### User Integrations as Destinations
Users configure their own upload targets via the **Integrations** dashboard
(`/integrations`). A DESTINATION integration stores:
- **Config** (`config` column, JSON): non-sensitive settings such as bucket
name, remote folder, SMTP host, etc.
- **Credentials** (`credentials` column, Fernet-encrypted JSON): sensitive
values such as OAuth refresh tokens, API keys, and passwords.
When uploading, credentials are decrypted at task execution time and passed
directly to the appropriate upload handler — they never appear in plain text
in task messages or logs.
### Supported Destination Types
| Integration Type | Upload Method |
|-----------------|--------------|
| `DROPBOX` | Dropbox SDK, OAuth refresh-token flow |
| `S3` | boto3 `upload_file`, per-user access key |
| `GOOGLE_DRIVE` | Google Drive API v3, OAuth or service account |
| `ONEDRIVE` | Microsoft Graph API, MSAL confidential-client |
| `WEBDAV` | HTTP PUT request, Basic Auth |
| `NEXTCLOUD` | WebDAV (same as WEBDAV, Nextcloud-compatible path) |
| `FTP` | ftplib FTPS (TLS preferred, plaintext configurable) |
| `SFTP` | Paramiko, password or private-key auth |
| `PAPERLESS` | Paperless-ngx REST API, API token |
| `EMAIL` | SMTP/STARTTLS, file as attachment |
| `RCLONE` | `rclone copyto` subprocess, per-user rclone config |
### Multiple Destinations
If a user configures multiple active DESTINATION integrations, the file is
uploaded to **each one asynchronously and independently**. Success or failure
per destination is logged separately so a single failing destination does not
block the others.
### Fallback to Global Destinations
Global destinations (configured via environment variables / admin settings)
are used whenever:
- The document has no owner (`owner_id` is `None`), e.g., uploaded in
single-user / anonymous mode.
- The owner exists but has **zero** active DESTINATION integrations.
This ensures backward compatibility with existing single-user deployments.
## See Also
- [API Documentation](API.md) - API endpoints for file operations
+236 -42
View File
@@ -6,13 +6,34 @@ import pytest
from app.tasks.finalize_document_storage import finalize_document_storage
# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------
def _make_file_record(file_id: int = 123, owner_id=None):
"""Return a lightweight MagicMock that mimics a FileRecord."""
rec = MagicMock()
rec.id = file_id
rec.owner_id = owner_id
return rec
@pytest.fixture(autouse=True)
def _patch_celery_background_tasks(mocker):
"""Module-level autouse fixture: prevent lazy-imported Celery tasks from connecting to Redis."""
mocker.patch("app.tasks.compute_embedding.compute_document_embedding")
mocker.patch("app.tasks.convert_to_pdfa.convert_to_pdfa", create=True)
@pytest.mark.unit
class TestFinalizeDocumentStorage:
"""Tests for finalize_document_storage Celery task."""
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
@@ -21,11 +42,12 @@ class TestFinalizeDocumentStorage:
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""Test successful document finalization with all services configured."""
# Mock configured services
mock_get_services.return_value = {
"dropbox": True,
"google_drive": True,
@@ -33,14 +55,11 @@ class TestFinalizeDocumentStorage:
"s3": True,
}
# Mock database session
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_file_record = MagicMock()
mock_file_record.id = 123
mock_file_record = _make_file_record(123, owner_id=None)
mock_db.query.return_value.filter.return_value.first.return_value = mock_file_record
# Mock file existence and size
with patch("app.tasks.finalize_document_storage.os.path.exists", return_value=True):
with patch("app.tasks.finalize_document_storage.os.path.getsize", return_value=102400):
with patch("app.tasks.finalize_document_storage.os.path.basename", return_value="test_document.pdf"):
@@ -59,10 +78,10 @@ class TestFinalizeDocumentStorage:
file_id=123,
)
# Verify send_to_all_destinations was queued
# owner_id=None → global routing
mock_send_all.delay.assert_called_once_with("/workdir/processed/test_document.pdf", True, 123)
mock_send_user.delay.assert_not_called()
# Verify notification was sent
mock_notify.assert_called_once()
notify_args = mock_notify.call_args[1]
assert notify_args["filename"] == "test_document.pdf"
@@ -72,12 +91,13 @@ class TestFinalizeDocumentStorage:
assert "Google Drive" in notify_args["destinations"]
assert "S3" in notify_args["destinations"]
# Verify result
assert result["status"] == "Completed"
assert result["file"] == "/workdir/processed/test_document.pdf"
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
@@ -86,17 +106,17 @@ class TestFinalizeDocumentStorage:
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""Test file_id retrieval from database when not provided."""
mock_get_services.return_value = {"dropbox": True}
# Mock database session to return a file record
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_file_record = MagicMock()
mock_file_record.id = 456
mock_file_record = _make_file_record(456, owner_id=None)
mock_db.query.return_value.filter.return_value.first.return_value = mock_file_record
with patch("app.tasks.finalize_document_storage.os.path.exists", return_value=True):
@@ -121,11 +141,13 @@ class TestFinalizeDocumentStorage:
# Verify database was queried
mock_db.query.assert_called_once()
# Verify send_to_all was called with retrieved file_id
# Verify send_to_all was called (global routing — no user destinations)
mock_send_all.delay.assert_called_once()
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
@@ -134,11 +156,12 @@ class TestFinalizeDocumentStorage:
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""Test handles case when no services are configured."""
# No services configured
mock_get_services.return_value = {
"dropbox": False,
"google_drive": False,
@@ -162,17 +185,17 @@ class TestFinalizeDocumentStorage:
file_id=789,
)
# Should still queue uploads (even if none configured)
# Should still queue global uploads (even if none configured)
mock_send_all.delay.assert_called_once()
# Should still send notification
mock_notify.assert_called_once()
notify_args = mock_notify.call_args[1]
# No services configured means empty destinations list
assert notify_args["destinations"] == []
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
@@ -181,11 +204,12 @@ class TestFinalizeDocumentStorage:
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""Test handles exception when getting configured services."""
# Simulate exception
mock_get_services.side_effect = Exception("Service validation failed")
mock_db = MagicMock()
@@ -204,16 +228,16 @@ class TestFinalizeDocumentStorage:
file_id=101,
)
# Should still complete successfully
assert result["status"] == "Completed"
# Should use fallback destinations
mock_notify.assert_called_once()
notify_args = mock_notify.call_args[1]
assert "configured destinations" in notify_args["destinations"]
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
@@ -222,7 +246,9 @@ class TestFinalizeDocumentStorage:
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""Test handles notification failure gracefully."""
@@ -232,7 +258,6 @@ class TestFinalizeDocumentStorage:
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = None
# Simulate notification failure
mock_notify.side_effect = Exception("Notification service unavailable")
with patch("app.tasks.finalize_document_storage.os.path.exists", return_value=True):
@@ -247,14 +272,13 @@ class TestFinalizeDocumentStorage:
file_id=202,
)
# Should still complete successfully despite notification failure
assert result["status"] == "Completed"
# Should still queue uploads
mock_send_all.delay.assert_called_once()
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
@@ -263,7 +287,9 @@ class TestFinalizeDocumentStorage:
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""Test handles case when processed file doesn't exist."""
@@ -273,7 +299,6 @@ class TestFinalizeDocumentStorage:
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = None
# File doesn't exist
with patch("app.tasks.finalize_document_storage.os.path.exists", return_value=False):
with patch("app.tasks.finalize_document_storage.os.path.basename", return_value="missing.pdf"):
finalize_document_storage.request.id = "test-task-id"
@@ -285,16 +310,16 @@ class TestFinalizeDocumentStorage:
file_id=303,
)
# Should still queue uploads (send_to_all handles missing files)
mock_send_all.delay.assert_called_once()
# Notification should use file_size = 0
mock_notify.assert_called_once()
notify_args = mock_notify.call_args[1]
assert notify_args["file_size"] == 0
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
@@ -303,11 +328,12 @@ class TestFinalizeDocumentStorage:
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""Test that service names are formatted correctly for display."""
# Mock services with underscores in names
mock_get_services.return_value = {
"google_drive": True,
"one_drive": True,
@@ -330,16 +356,17 @@ class TestFinalizeDocumentStorage:
file_id=404,
)
# Verify service names are formatted with spaces and title case
mock_notify.assert_called_once()
notify_args = mock_notify.call_args[1]
destinations = notify_args["destinations"]
assert "Google Drive" in destinations
assert "One Drive" in destinations
assert "Next Cloud" not in destinations # Not configured
assert "Next Cloud" not in destinations
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
@@ -348,7 +375,9 @@ class TestFinalizeDocumentStorage:
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""Test that delete_after flag is correctly passed to send_to_all_destinations."""
@@ -370,11 +399,12 @@ class TestFinalizeDocumentStorage:
file_id=505,
)
# Verify send_to_all was called with delete_after=True
mock_send_all.delay.assert_called_once_with("/workdir/processed/file.pdf", True, 505)
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
@@ -383,17 +413,14 @@ class TestFinalizeDocumentStorage:
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""
Regression test: when PDF/A conversion is enabled, the finalize_document_storage
step must NOT be logged as in_progress after it has already been logged as success.
Previously, a second log_task_progress call with status="in_progress" was made for
"finalize_document_storage" when queueing PDF/A archival conversion, which overwrote
the prior success status and caused the overall file status to appear stuck in
processing/failed.
"""
mock_get_services.return_value = {"dropbox": True}
@@ -423,18 +450,185 @@ class TestFinalizeDocumentStorage:
file_id=606,
)
# Collect all (step_name, status) pairs logged for finalize_document_storage
# Collect all logged calls for finalize_document_storage step
finalize_calls = [
call
for call in mock_log_progress.call_args_list
if call.args[1] == "finalize_document_storage"
c for c in mock_log_progress.call_args_list if c.args[1] == "finalize_document_storage"
]
# After the success log, no in_progress log should follow for this step
statuses = [call.args[2] for call in finalize_calls]
statuses = [c.args[2] for c in finalize_calls]
assert "success" in statuses, "finalize_document_storage must be logged as success"
# The last status logged must be success, not in_progress
assert statuses[-1] == "success", (
"finalize_document_storage must not be regressed to in_progress after success; "
f"got statuses: {statuses}"
)
@pytest.mark.unit
class TestFinalizeDocumentStorageUserRouting:
"""Tests for user-specific destination routing in finalize_document_storage."""
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=2)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
def test_routes_to_user_destinations_when_owner_has_integrations(
self,
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""When a user has active DESTINATION integrations, use them instead of global config."""
mock_get_services.return_value = {"dropbox": True}
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_file_record = _make_file_record(100, owner_id="alice@example.com")
mock_db.query.return_value.filter.return_value.first.return_value = mock_file_record
with patch("app.tasks.finalize_document_storage.os.path.exists", return_value=True):
with patch("app.tasks.finalize_document_storage.os.path.getsize", return_value=1024):
with patch("app.tasks.finalize_document_storage.os.path.basename", return_value="doc.pdf"):
finalize_document_storage.request.id = "test-task-id"
result = finalize_document_storage.__wrapped__(
original_file="/tmp/original.pdf",
processed_file="/workdir/processed/doc.pdf",
metadata={"filename": "doc.pdf"},
file_id=100,
)
mock_send_user.delay.assert_called_once_with("/workdir/processed/doc.pdf", "alice@example.com", 100)
mock_send_all.delay.assert_not_called()
assert result["status"] == "Completed"
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
def test_falls_back_to_global_when_owner_has_no_integrations(
self,
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""When a user has no active DESTINATION integrations, fall back to global config."""
mock_get_services.return_value = {"s3": True}
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_file_record = _make_file_record(200, owner_id="bob@example.com")
mock_db.query.return_value.filter.return_value.first.return_value = mock_file_record
with patch("app.tasks.finalize_document_storage.os.path.exists", return_value=True):
with patch("app.tasks.finalize_document_storage.os.path.getsize", return_value=2048):
with patch("app.tasks.finalize_document_storage.os.path.basename", return_value="file.pdf"):
finalize_document_storage.request.id = "test-task-id"
result = finalize_document_storage.__wrapped__(
original_file="/tmp/original.pdf",
processed_file="/workdir/processed/file.pdf",
metadata={"filename": "file.pdf"},
file_id=200,
)
mock_send_all.delay.assert_called_once_with("/workdir/processed/file.pdf", True, 200)
mock_send_user.delay.assert_not_called()
assert result["status"] == "Completed"
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
def test_falls_back_to_global_when_no_owner(
self,
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""When a document has no owner (single-user mode), global destinations are used."""
mock_get_services.return_value = {"nextcloud": True}
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_file_record = _make_file_record(300, owner_id=None)
mock_db.query.return_value.filter.return_value.first.return_value = mock_file_record
with patch("app.tasks.finalize_document_storage.os.path.exists", return_value=True):
with patch("app.tasks.finalize_document_storage.os.path.getsize", return_value=512):
with patch("app.tasks.finalize_document_storage.os.path.basename", return_value="scan.pdf"):
finalize_document_storage.request.id = "test-task-id"
result = finalize_document_storage.__wrapped__(
original_file="/tmp/original.pdf",
processed_file="/workdir/processed/scan.pdf",
metadata={"filename": "scan.pdf"},
file_id=300,
)
mock_send_all.delay.assert_called_once_with("/workdir/processed/scan.pdf", True, 300)
mock_send_user.delay.assert_not_called()
# get_user_destination_count must NOT be called when owner_id is None
mock_get_dest_count.assert_not_called()
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count")
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
def test_falls_back_to_global_when_count_lookup_fails(
self,
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
):
"""When get_user_destination_count raises, fall back to global routing gracefully."""
mock_get_services.return_value = {"s3": True}
mock_get_dest_count.side_effect = Exception("DB connection error")
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_file_record = _make_file_record(400, owner_id="charlie@example.com")
mock_db.query.return_value.filter.return_value.first.return_value = mock_file_record
with patch("app.tasks.finalize_document_storage.os.path.exists", return_value=True):
with patch("app.tasks.finalize_document_storage.os.path.getsize", return_value=4096):
with patch("app.tasks.finalize_document_storage.os.path.basename", return_value="file.pdf"):
finalize_document_storage.request.id = "test-task-id"
result = finalize_document_storage.__wrapped__(
original_file="/tmp/original.pdf",
processed_file="/workdir/processed/file.pdf",
metadata={"filename": "file.pdf"},
file_id=400,
)
mock_send_all.delay.assert_called_once_with("/workdir/processed/file.pdf", True, 400)
mock_send_user.delay.assert_not_called()
assert result["status"] == "Completed"
+958
View File
@@ -0,0 +1,958 @@
"""Unit tests for the per-type upload handler functions in upload_to_user_integration.py.
Each ``_upload_*`` helper is tested by mocking the relevant third-party library
so that tests are fast, hermetic, and free of external network calls.
"""
import json
import os
from unittest.mock import MagicMock, patch
import pytest
# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------
TASK_ID = "test-handler-task-id"
def _write_file(path, content: bytes = b"PDF content") -> None:
"""Write *content* to *path*, creating parent dirs as needed."""
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "wb") as fh:
fh.write(content)
# ---------------------------------------------------------------------------
# _upload_dropbox
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadDropbox:
"""Tests for _upload_dropbox handler."""
def _call(self, file_path: str, cfg: dict, creds: dict) -> dict:
from app.tasks.upload_to_user_integration import _upload_dropbox
return _upload_dropbox(file_path, cfg, creds, TASK_ID)
def test_raises_when_missing_credentials(self, tmp_path):
"""ValueError raised when app_key, app_secret, or refresh_token is missing."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="app_key"):
self._call(fp, {}, {})
def test_small_file_upload(self, tmp_path):
"""Files ≤10 MB are uploaded with files_upload."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp, b"x" * 100)
mock_dbx_instance = MagicMock()
mock_dropbox_files = MagicMock()
mock_dropbox_files.WriteMode.overwrite = "overwrite"
with patch.dict(
"sys.modules",
{
"dropbox": MagicMock(
Dropbox=MagicMock(return_value=mock_dbx_instance),
files=mock_dropbox_files,
)
},
):
result = self._call(
fp,
{"folder": "/Docs"},
{"app_key": "key", "app_secret": "secret", "refresh_token": "rtoken"},
)
mock_dbx_instance.files_upload.assert_called_once()
assert result["status"] == "Completed"
assert result["dropbox_path"] == "/Docs/doc.pdf"
def test_large_file_upload_uses_session(self, tmp_path):
"""Files >10 MB are uploaded with upload session (chunked)."""
fp = str(tmp_path / "large.pdf")
# Write 11 MB
_write_file(fp, b"x" * (11 * 1024 * 1024))
mock_dbx_instance = MagicMock()
mock_session_start = MagicMock()
mock_session_start.session_id = "session-1"
mock_dbx_instance.files_upload_session_start.return_value = mock_session_start
mock_dbx_instance.files_upload_session_finish.return_value = MagicMock()
mock_files_mod = MagicMock()
mock_files_mod.UploadSessionCursor = MagicMock(return_value=MagicMock(offset=0))
mock_files_mod.CommitInfo = MagicMock()
mock_files_mod.WriteMode.overwrite = "overwrite"
with patch.dict(
"sys.modules",
{
"dropbox": MagicMock(
Dropbox=MagicMock(return_value=mock_dbx_instance),
files=mock_files_mod,
)
},
):
result = self._call(
fp,
{},
{"app_key": "k", "app_secret": "s", "refresh_token": "r"},
)
mock_dbx_instance.files_upload_session_start.assert_called_once()
assert result["status"] == "Completed"
def test_default_folder_when_not_specified(self, tmp_path):
"""When no folder is configured, the default '/DocuElevate' folder is used."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp, b"x" * 10)
mock_dbx_instance = MagicMock()
with patch.dict(
"sys.modules",
{
"dropbox": MagicMock(
Dropbox=MagicMock(return_value=mock_dbx_instance),
files=MagicMock(WriteMode=MagicMock(overwrite="overwrite")),
)
},
):
result = self._call(fp, {}, {"app_key": "k", "app_secret": "s", "refresh_token": "r"})
assert result["dropbox_path"] == "/DocuElevate/doc.pdf"
# ---------------------------------------------------------------------------
# _upload_s3
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadS3:
"""Tests for _upload_s3 handler."""
def _call(self, file_path: str, cfg: dict, creds: dict) -> dict:
from app.tasks.upload_to_user_integration import _upload_s3
return _upload_s3(file_path, cfg, creds, TASK_ID)
def test_raises_when_bucket_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="bucket"):
self._call(fp, {}, {"access_key_id": "k", "secret_access_key": "s"})
def test_raises_when_credentials_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="access_key_id"):
self._call(fp, {"bucket": "my-bucket"}, {})
def test_successful_upload(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_s3 = MagicMock()
mock_boto3 = MagicMock()
mock_boto3.client.return_value = mock_s3
with patch.dict("sys.modules", {"boto3": mock_boto3, "botocore.exceptions": MagicMock(ClientError=Exception)}):
result = self._call(
fp,
{"bucket": "my-bucket", "region": "eu-west-1", "folder_prefix": "docs"},
{"access_key_id": "AKIA", "secret_access_key": "secret"},
)
mock_s3.upload_file.assert_called_once()
assert result["status"] == "Completed"
assert result["s3_bucket"] == "my-bucket"
assert result["s3_key"] == "docs/doc.pdf"
def test_uses_endpoint_url_when_provided(self, tmp_path):
"""Custom endpoint_url is passed to boto3.client for S3-compatible stores."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_s3 = MagicMock()
mock_boto3 = MagicMock()
mock_boto3.client.return_value = mock_s3
with patch.dict("sys.modules", {"boto3": mock_boto3, "botocore.exceptions": MagicMock(ClientError=Exception)}):
self._call(
fp,
{"bucket": "b", "endpoint_url": "https://minio.example.com"},
{"access_key_id": "k", "secret_access_key": "s"},
)
call_kwargs = mock_boto3.client.call_args[1]
assert call_kwargs.get("endpoint_url") == "https://minio.example.com"
def test_wraps_client_error_as_runtime_error(self, tmp_path):
"""S3 ClientError is re-raised as RuntimeError."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
class FakeClientError(Exception):
pass
mock_s3 = MagicMock()
mock_s3.upload_file.side_effect = FakeClientError("Access Denied")
mock_boto3 = MagicMock()
mock_boto3.client.return_value = mock_s3
with patch.dict(
"sys.modules",
{"boto3": mock_boto3, "botocore.exceptions": MagicMock(ClientError=FakeClientError)},
):
with pytest.raises(RuntimeError, match="S3 upload failed"):
self._call(fp, {"bucket": "b"}, {"access_key_id": "k", "secret_access_key": "s"})
def test_key_without_folder_prefix(self, tmp_path):
"""When folder_prefix is empty, the S3 key is just the filename."""
fp = str(tmp_path / "report.pdf")
_write_file(fp)
mock_s3 = MagicMock()
mock_boto3 = MagicMock()
mock_boto3.client.return_value = mock_s3
with patch.dict("sys.modules", {"boto3": mock_boto3, "botocore.exceptions": MagicMock(ClientError=Exception)}):
result = self._call(fp, {"bucket": "b"}, {"access_key_id": "k", "secret_access_key": "s"})
assert result["s3_key"] == "report.pdf"
# ---------------------------------------------------------------------------
# _upload_google_drive
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadGoogleDrive:
"""Tests for _upload_google_drive handler."""
def _call(self, file_path: str, cfg: dict, creds: dict) -> dict:
from app.tasks.upload_to_user_integration import _upload_google_drive
return _upload_google_drive(file_path, cfg, creds, TASK_ID)
def test_raises_when_no_credentials(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="OAuth credentials"):
self._call(fp, {}, {})
def test_oauth_upload_calls_drive_api(self, tmp_path):
"""OAuth credentials (client_id + client_secret + refresh_token) trigger OAuth flow."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_service = MagicMock()
mock_service.files.return_value.create.return_value.execute.return_value = {
"id": "gdrive-id-123",
"webViewLink": "https://drive.google.com/file/d/gdrive-id-123",
}
mock_build = MagicMock(return_value=mock_service)
mock_oauth_creds = MagicMock()
mock_google_oauth2 = MagicMock()
mock_google_oauth2.credentials.Credentials = MagicMock(return_value=mock_oauth_creds)
mock_google_auth_transport = MagicMock()
mock_google_auth_transport.requests.Request = MagicMock()
mock_media_upload = MagicMock()
with patch.dict(
"sys.modules",
{
"googleapiclient.discovery": MagicMock(build=mock_build),
"googleapiclient.http": MagicMock(MediaFileUpload=mock_media_upload),
"google.oauth2.credentials": mock_google_oauth2.credentials,
"google.auth.transport.requests": mock_google_auth_transport.requests,
"google.oauth2.service_account": MagicMock(),
},
):
result = self._call(
fp,
{"folder_id": "folder-xyz"},
{"client_id": "cid", "client_secret": "csec", "refresh_token": "rtoken"},
)
assert result["status"] == "Completed"
assert result["google_drive_file_id"] == "gdrive-id-123"
def test_service_account_upload(self, tmp_path):
"""credentials_json triggers service-account flow."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
sa_creds_json = json.dumps({"type": "service_account", "project_id": "myproject"})
mock_service = MagicMock()
mock_service.files.return_value.create.return_value.execute.return_value = {
"id": "sa-file-id",
"webViewLink": "https://drive.google.com/file/d/sa-file-id",
}
mock_sa_class = MagicMock()
mock_sa_creds = MagicMock()
mock_sa_class.from_service_account_info.return_value = mock_sa_creds
mock_build = MagicMock(return_value=mock_service)
mock_media_upload = MagicMock()
with patch.dict(
"sys.modules",
{
"googleapiclient.discovery": MagicMock(build=mock_build),
"googleapiclient.http": MagicMock(MediaFileUpload=mock_media_upload),
"google.oauth2.credentials": MagicMock(),
"google.auth.transport.requests": MagicMock(),
"google.oauth2.service_account": MagicMock(Credentials=mock_sa_class),
},
):
result = self._call(fp, {}, {"credentials_json": sa_creds_json})
assert result["status"] == "Completed"
# ---------------------------------------------------------------------------
# _upload_webdav / _upload_nextcloud
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadWebdav:
"""Tests for _upload_webdav handler (and Nextcloud which delegates to it)."""
def _call(self, file_path: str, cfg: dict, creds: dict) -> dict:
from app.tasks.upload_to_user_integration import _upload_webdav
return _upload_webdav(file_path, cfg, creds, TASK_ID)
def test_raises_when_url_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="url"):
self._call(fp, {}, {})
def test_successful_upload_201(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_resp = MagicMock()
mock_resp.status_code = 201
mock_requests = MagicMock()
mock_requests.put.return_value = mock_resp
with patch.dict("sys.modules", {"requests": mock_requests}):
result = self._call(
fp,
{"url": "https://dav.example.com/dav/", "folder": "Files"},
{"username": "user", "password": "pass"},
)
assert result["status"] == "Completed"
mock_requests.put.assert_called_once()
def test_raises_on_non_2xx_response(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_resp = MagicMock()
mock_resp.status_code = 403
mock_resp.text = "Forbidden"
mock_requests = MagicMock()
mock_requests.put.return_value = mock_resp
with patch.dict("sys.modules", {"requests": mock_requests}):
with pytest.raises(RuntimeError, match="WebDAV upload failed: 403"):
self._call(fp, {"url": "https://dav.example.com/"}, {})
def test_nextcloud_delegates_to_webdav(self, tmp_path):
"""_upload_nextcloud is a thin wrapper over _upload_webdav."""
from app.tasks.upload_to_user_integration import _upload_nextcloud
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with patch("app.tasks.upload_to_user_integration._upload_webdav") as mock_webdav:
mock_webdav.return_value = {"status": "Completed", "webdav_url": "https://nc.example.com/Files/doc.pdf"}
result = _upload_nextcloud(
fp, {"url": "https://nc.example.com"}, {"username": "u", "password": "p"}, TASK_ID
)
mock_webdav.assert_called_once_with(
fp, {"url": "https://nc.example.com"}, {"username": "u", "password": "p"}, TASK_ID
)
assert result["status"] == "Completed"
# ---------------------------------------------------------------------------
# _upload_ftp
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadFtp:
"""Tests for _upload_ftp handler."""
def _call(self, file_path: str, cfg: dict, creds: dict) -> dict:
from app.tasks.upload_to_user_integration import _upload_ftp
return _upload_ftp(file_path, cfg, creds, TASK_ID)
def test_raises_when_host_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="host"):
self._call(fp, {}, {"password": "pass"})
def test_tls_upload(self, tmp_path):
"""use_tls=True uses FTP_TLS."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_ftp_tls = MagicMock()
mock_ftplib = MagicMock()
mock_ftplib.FTP_TLS.return_value = mock_ftp_tls
mock_ftplib.error_perm = Exception
with patch.dict("sys.modules", {"ftplib": mock_ftplib}):
with patch("app.tasks.upload_to_user_integration.ftplib", mock_ftplib):
result = self._call(
fp,
{"host": "ftp.example.com", "port": 21, "folder": "/docs", "use_tls": True},
{"password": "pass"},
)
mock_ftplib.FTP_TLS.assert_called_once()
assert result["status"] == "Completed"
def test_plaintext_ftp_upload(self, tmp_path):
"""use_tls=False uses plain FTP."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_ftp = MagicMock()
mock_ftplib = MagicMock()
mock_ftplib.FTP.return_value = mock_ftp
mock_ftplib.error_perm = Exception
with patch("app.tasks.upload_to_user_integration.ftplib", mock_ftplib):
result = self._call(
fp,
{"host": "ftp.example.com", "use_tls": False},
{"password": "pass"},
)
mock_ftplib.FTP.assert_called_once()
assert result["status"] == "Completed"
def test_creates_folder_if_cwd_fails(self, tmp_path):
"""When cwd raises error_perm, the handler creates the directory."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
class FtpPermError(Exception):
pass
# cwd call sequence:
# 1. ftp.cwd("uploads") → fails (outer try, folder_stripped="uploads")
# 2. ftp.cwd("/uploads") → fails (inner loop, triggers mkd)
# 3. ftp.cwd("/uploads") after mkd → succeeds
mock_ftp = MagicMock()
mock_ftp.cwd.side_effect = [FtpPermError("no"), FtpPermError("no"), None]
mock_ftplib = MagicMock()
mock_ftplib.FTP.return_value = mock_ftp
mock_ftplib.error_perm = FtpPermError
with patch("app.tasks.upload_to_user_integration.ftplib", mock_ftplib):
self._call(
fp,
{"host": "ftp.example.com", "folder": "/uploads", "use_tls": False},
{"password": "p"},
)
mock_ftp.mkd.assert_called_with("/uploads")
# ---------------------------------------------------------------------------
# _upload_sftp
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadSftp:
"""Tests for _upload_sftp handler."""
def _call(self, file_path: str, cfg: dict, creds: dict) -> dict:
from app.tasks.upload_to_user_integration import _upload_sftp
return _upload_sftp(file_path, cfg, creds, TASK_ID)
def test_raises_when_host_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="host"):
self._call(fp, {}, {"password": "p"})
def test_raises_when_no_auth(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="password or private_key"):
self._call(fp, {"host": "sftp.example.com"}, {})
def test_password_auth(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_ssh = MagicMock()
mock_sftp = MagicMock()
mock_ssh.open_sftp.return_value = mock_sftp
mock_paramiko = MagicMock()
mock_paramiko.SSHClient.return_value = mock_ssh
mock_paramiko.RejectPolicy = MagicMock
with patch.dict("sys.modules", {"paramiko": mock_paramiko}):
result = self._call(
fp,
{"host": "sftp.example.com", "username": "user", "folder": "/uploads"},
{"password": "pass"},
)
mock_sftp.put.assert_called_once()
assert result["status"] == "Completed"
assert result["sftp_host"] == "sftp.example.com"
def test_private_key_auth(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_ssh = MagicMock()
mock_sftp = MagicMock()
mock_ssh.open_sftp.return_value = mock_sftp
mock_pkey = MagicMock()
mock_paramiko = MagicMock()
mock_paramiko.SSHClient.return_value = mock_ssh
mock_paramiko.RSAKey.from_private_key.return_value = mock_pkey
mock_paramiko.RejectPolicy = MagicMock
with patch.dict("sys.modules", {"paramiko": mock_paramiko}):
result = self._call(
fp,
{"host": "sftp.example.com", "username": "user"},
{"private_key": "-----BEGIN RSA PRIVATE KEY-----\n...\n-----END RSA PRIVATE KEY-----"},
)
mock_paramiko.RSAKey.from_private_key.assert_called_once()
assert result["status"] == "Completed"
# ---------------------------------------------------------------------------
# _upload_paperless
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadPaperless:
"""Tests for _upload_paperless handler."""
def _call(self, file_path: str, cfg: dict, creds: dict) -> dict:
from app.tasks.upload_to_user_integration import _upload_paperless
return _upload_paperless(file_path, cfg, creds, TASK_ID)
def test_raises_when_host_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="host"):
self._call(fp, {}, {"api_token": "tok"})
def test_raises_when_api_token_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="api_token"):
self._call(fp, {"host": "https://paperless.example.com"}, {})
def test_successful_upload_polls_to_success(self, tmp_path):
"""Document is uploaded and task polling returns SUCCESS."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
# POST response
mock_post_resp = MagicMock()
mock_post_resp.text = '"task-uuid-123"'
# Poll response showing SUCCESS
mock_poll_resp = MagicMock()
mock_poll_resp.json.return_value = [{"status": "SUCCESS", "related_document": 42}]
mock_requests = MagicMock()
mock_requests.post.return_value = mock_post_resp
mock_requests.get.return_value = mock_poll_resp
with patch.dict("sys.modules", {"requests": mock_requests}):
with patch("time.sleep", return_value=None):
result = self._call(
fp,
{"host": "https://paperless.example.com"},
{"api_token": "tok-abc"},
)
assert result["status"] == "Completed"
assert result["paperless_document_id"] == 42
def test_raises_when_paperless_task_fails(self, tmp_path):
"""RuntimeError is raised when Paperless processing status is FAILURE."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_post_resp = MagicMock()
mock_post_resp.text = '"task-uuid-999"'
mock_poll_resp = MagicMock()
mock_poll_resp.json.return_value = [{"status": "FAILURE", "result": "OCR failed"}]
mock_requests = MagicMock()
mock_requests.post.return_value = mock_post_resp
mock_requests.get.return_value = mock_poll_resp
with patch.dict("sys.modules", {"requests": mock_requests}):
with patch("time.sleep", return_value=None):
with pytest.raises(RuntimeError, match="Paperless processing failed"):
self._call(
fp,
{"host": "https://paperless.example.com"},
{"api_token": "tok"},
)
# ---------------------------------------------------------------------------
# _upload_email
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadEmail:
"""Tests for _upload_email handler."""
def _call(self, file_path: str, cfg: dict, creds: dict) -> dict:
from app.tasks.upload_to_user_integration import _upload_email
return _upload_email(file_path, cfg, creds, TASK_ID)
def test_raises_when_host_or_recipient_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="host or recipient"):
self._call(fp, {}, {})
def test_tls_email_sent(self, tmp_path):
"""use_tls=True invokes starttls() with ssl context."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_smtp_instance = MagicMock()
mock_smtp_class = MagicMock(return_value=mock_smtp_instance)
mock_smtp_instance.__enter__ = MagicMock(return_value=mock_smtp_instance)
mock_smtp_instance.__exit__ = MagicMock(return_value=False)
mock_ssl_ctx = MagicMock()
mock_ssl = MagicMock()
mock_ssl.create_default_context.return_value = mock_ssl_ctx
with patch("smtplib.SMTP", mock_smtp_class):
with patch("ssl.create_default_context", return_value=mock_ssl_ctx):
result = self._call(
fp,
{
"host": "smtp.example.com",
"port": 587,
"username": "u@ex.com",
"recipient": "r@ex.com",
"use_tls": True,
},
{"password": "pass"},
)
mock_smtp_instance.starttls.assert_called_once_with(context=mock_ssl_ctx)
assert result["status"] == "Completed"
assert result["recipient"] == "r@ex.com"
def test_plaintext_smtp_skips_starttls(self, tmp_path):
"""use_tls=False sends without starttls()."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_smtp_instance = MagicMock()
mock_smtp_class = MagicMock(return_value=mock_smtp_instance)
mock_smtp_instance.__enter__ = MagicMock(return_value=mock_smtp_instance)
mock_smtp_instance.__exit__ = MagicMock(return_value=False)
with patch("smtplib.SMTP", mock_smtp_class):
result = self._call(
fp,
{"host": "smtp.example.com", "recipient": "r@ex.com", "use_tls": False},
{},
)
mock_smtp_instance.starttls.assert_not_called()
assert result["status"] == "Completed"
# ---------------------------------------------------------------------------
# _upload_rclone
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadRclone:
"""Tests for _upload_rclone handler."""
def _call(self, file_path: str, cfg: dict, creds: dict) -> dict:
from app.tasks.upload_to_user_integration import _upload_rclone
return _upload_rclone(file_path, cfg, creds, TASK_ID)
def test_raises_when_remote_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="remote"):
self._call(fp, {}, {"rclone_conf": "[myremote]\ntype = s3\n"})
def test_raises_when_conf_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="rclone_conf"):
self._call(fp, {"remote": "myremote:"}, {})
def test_raises_when_remote_unsafe(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="unsafe characters"):
self._call(fp, {"remote": "my;remote:"}, {"rclone_conf": "[x]\ntype=s3\n"})
def test_raises_when_folder_unsafe(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="unsafe characters"):
self._call(fp, {"remote": "myremote:", "folder": "docs;rm -rf /"}, {"rclone_conf": "[x]\ntype=s3\n"})
def test_successful_rclone_copy(self, tmp_path):
"""rclone process is called with correct arguments and temp config file."""
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_result = MagicMock()
mock_result.returncode = 0
with patch("app.tasks.upload_to_user_integration.subprocess.run", return_value=mock_result) as mock_run:
result = self._call(
fp,
{"remote": "myremote:", "folder": "docs"},
{"rclone_conf": "[myremote]\ntype = s3\n"},
)
assert result["status"] == "Completed"
# Verify subprocess.run was called with rclone command
cmd = mock_run.call_args[0][0]
assert cmd[0] == "rclone"
assert cmd[1] == "copyto"
def test_raises_on_rclone_nonzero_exit(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_result = MagicMock()
mock_result.returncode = 1
mock_result.stderr = "rclone: command not found"
with patch("app.tasks.upload_to_user_integration.subprocess.run", return_value=mock_result):
with pytest.raises(RuntimeError, match="rclone exited 1"):
self._call(fp, {"remote": "myremote:"}, {"rclone_conf": "[myremote]\ntype=s3\n"})
# ---------------------------------------------------------------------------
# _upload_onedrive
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadOneDrive:
"""Tests for _upload_onedrive handler."""
def _call(self, file_path: str, cfg: dict, creds: dict) -> dict:
from app.tasks.upload_to_user_integration import _upload_onedrive
return _upload_onedrive(file_path, cfg, creds, TASK_ID)
def test_raises_when_client_credentials_missing(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
with pytest.raises(ValueError, match="client_id or client_secret"):
self._call(fp, {}, {})
def test_raises_when_token_acquisition_fails(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp)
mock_msal_app = MagicMock()
mock_msal_app.acquire_token_for_client.return_value = {
"error": "invalid_client",
"error_description": "AADSTS70011",
}
mock_msal = MagicMock()
mock_msal.ConfidentialClientApplication.return_value = mock_msal_app
with patch.dict("sys.modules", {"msal": mock_msal, "requests": MagicMock()}):
with pytest.raises(ValueError, match="token acquisition failed"):
self._call(fp, {}, {"client_id": "cid", "client_secret": "csec"})
def test_successful_upload_with_refresh_token(self, tmp_path):
fp = str(tmp_path / "doc.pdf")
_write_file(fp, b"x" * 100)
mock_msal_app = MagicMock()
mock_msal_app.acquire_token_by_refresh_token.return_value = {"access_token": "tok-abc"}
mock_msal = MagicMock()
mock_msal.ConfidentialClientApplication.return_value = mock_msal_app
# Mock POST (create upload session) and PUT (chunk upload)
mock_post_resp = MagicMock()
mock_post_resp.json.return_value = {"uploadUrl": "https://upload.example.com/session"}
mock_put_resp = MagicMock()
mock_put_resp.status_code = 201
mock_requests = MagicMock()
mock_requests.post.return_value = mock_post_resp
mock_requests.put.return_value = mock_put_resp
with patch.dict("sys.modules", {"msal": mock_msal, "requests": mock_requests}):
result = self._call(
fp,
{"folder_path": "Documents/DocuElevate", "tenant_id": "my-tenant"},
{"client_id": "cid", "client_secret": "csec", "refresh_token": "rtoken"},
)
assert result["status"] == "Completed"
# ---------------------------------------------------------------------------
# finalize_document_storage - uncovered branches
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestFinalizeDocumentStorageUncoveredBranches:
"""Cover branches in finalize_document_storage not exercised by the main test class."""
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count", return_value=0)
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
def test_resolves_owner_id_when_file_id_none(
self,
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
tmp_path,
):
"""When file_id is None, the task looks up the FileRecord by local_filename."""
from app.tasks.finalize_document_storage import finalize_document_storage
processed_file = str(tmp_path / "processed" / "doc.pdf")
original_file = str(tmp_path / "original" / "orig.pdf")
os.makedirs(os.path.dirname(processed_file), exist_ok=True)
os.makedirs(os.path.dirname(original_file), exist_ok=True)
_write_file(processed_file)
_write_file(original_file)
mock_file_record = MagicMock()
mock_file_record.id = 77
mock_file_record.owner_id = "owner@example.com"
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = mock_file_record
mock_get_services.return_value = {}
finalize_document_storage.request.id = "test-task-id"
finalize_document_storage.__wrapped__(
original_file=original_file,
processed_file=processed_file,
metadata={},
file_id=None,
)
# The function should have queried the DB for the file record
mock_db.query.assert_called()
@patch("app.tasks.finalize_document_storage.notify_file_processed")
@patch("app.tasks.finalize_document_storage.send_to_user_destinations")
@patch("app.tasks.finalize_document_storage.send_to_all_destinations")
@patch("app.tasks.finalize_document_storage.get_user_destination_count")
@patch("app.tasks.finalize_document_storage.get_configured_services_from_validator")
@patch("app.tasks.finalize_document_storage.log_task_progress")
@patch("app.tasks.finalize_document_storage.SessionLocal")
def test_routes_to_global_when_count_query_raises(
self,
mock_session_local,
mock_log_progress,
mock_get_services,
mock_get_dest_count,
mock_send_all,
mock_send_user,
mock_notify,
tmp_path,
):
"""When get_user_destination_count raises, falls back to global routing."""
from app.tasks.finalize_document_storage import finalize_document_storage
processed_file = str(tmp_path / "processed2" / "doc.pdf")
original_file = str(tmp_path / "original2" / "orig.pdf")
os.makedirs(os.path.dirname(processed_file), exist_ok=True)
os.makedirs(os.path.dirname(original_file), exist_ok=True)
_write_file(processed_file)
_write_file(original_file)
mock_file_record = MagicMock()
mock_file_record.id = 88
mock_file_record.owner_id = "owner@example.com"
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = mock_file_record
mock_get_services.return_value = {}
mock_get_dest_count.side_effect = Exception("DB connection error")
finalize_document_storage.request.id = "test-task-id"
finalize_document_storage.__wrapped__(
original_file=original_file,
processed_file=processed_file,
metadata={},
file_id=88,
)
# Falls back to global since count raised
mock_send_all.delay.assert_called_once()
mock_send_user.delay.assert_not_called()
+380
View File
@@ -0,0 +1,380 @@
"""Unit tests for app/tasks/upload_to_user_integration.py and related helpers."""
import json
from unittest.mock import MagicMock, patch
import pytest
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_integration(
int_id: int = 1,
int_type=None,
owner_id: str = "user@example.com",
name: str = "My Integration",
config_dict: dict | None = None,
creds_dict: dict | None = None,
):
"""Build a MagicMock resembling a UserIntegration row."""
from app.models import IntegrationDirection, IntegrationType
if int_type is None:
int_type = IntegrationType.S3
rec = MagicMock()
rec.id = int_id
rec.integration_type = int_type
rec.owner_id = owner_id
rec.name = name
rec.config = json.dumps(config_dict or {})
rec.credentials = json.dumps(creds_dict or {}) # plain JSON in tests (not encrypted)
rec.is_active = True
rec.direction = IntegrationDirection.DESTINATION
# Prevent last_used_at / last_error from being MagicMock initially
rec.last_used_at = None
rec.last_error = None
return rec
def _run_upload_task(file_path: str, integration_id: int, file_id: int | None = None):
"""Call the upload_to_user_integration task's __wrapped__ function directly."""
from app.tasks.upload_to_user_integration import upload_to_user_integration
upload_to_user_integration.request.id = "test-task-id"
return upload_to_user_integration.__wrapped__(
file_path=file_path,
integration_id=integration_id,
file_id=file_id,
)
# ---------------------------------------------------------------------------
# Tests for upload_to_user_integration task
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestUploadToUserIntegration:
"""Unit tests for the upload_to_user_integration Celery task."""
@patch("app.tasks.upload_to_user_integration.log_task_progress")
@patch("app.tasks.upload_to_user_integration.SessionLocal")
def test_raises_file_not_found(self, mock_session_local, mock_log_progress):
"""FileNotFoundError is raised when the file does not exist."""
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
with pytest.raises(FileNotFoundError):
_run_upload_task("/nonexistent/file.pdf", integration_id=1)
@patch("app.tasks.upload_to_user_integration.log_task_progress")
@patch("app.tasks.upload_to_user_integration.SessionLocal")
def test_raises_value_error_when_integration_not_found(self, mock_session_local, mock_log_progress, tmp_path):
"""ValueError is raised when the integration record does not exist."""
test_file = tmp_path / "doc.pdf"
test_file.write_bytes(b"PDF content")
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = None
with pytest.raises(ValueError, match="not found"):
_run_upload_task(str(test_file), integration_id=99)
@patch("app.tasks.upload_to_user_integration.log_task_progress")
@patch("app.tasks.upload_to_user_integration.SessionLocal")
def test_skips_imap_source_type(self, mock_session_local, mock_log_progress, tmp_path):
"""Integration types with no registered handler return status='Skipped'."""
test_file = tmp_path / "doc.pdf"
test_file.write_bytes(b"PDF content")
from app.models import IntegrationType
integration = _make_integration(int_id=2, int_type=IntegrationType.IMAP)
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = integration
result = _run_upload_task(str(test_file), integration_id=2)
assert result["status"] == "Skipped"
@patch("app.tasks.upload_to_user_integration.log_task_progress")
@patch("app.tasks.upload_to_user_integration.SessionLocal")
def test_dispatches_to_correct_handler(self, mock_session_local, mock_log_progress, tmp_path):
"""The correct handler is called for a given integration type."""
test_file = tmp_path / "doc.pdf"
test_file.write_bytes(b"PDF content")
from app.models import IntegrationType
from app.tasks.upload_to_user_integration import _UPLOAD_HANDLERS
integration = _make_integration(
int_id=3,
int_type=IntegrationType.S3,
config_dict={"bucket": "my-bucket", "region": "us-east-1"},
creds_dict={"access_key_id": "AKI...", "secret_access_key": "secret"},
)
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = integration
mock_handler = MagicMock(return_value={"status": "Completed", "s3_key": "doc.pdf"})
with patch.dict(_UPLOAD_HANDLERS, {IntegrationType.S3: mock_handler}):
result = _run_upload_task(str(test_file), integration_id=3, file_id=42)
mock_handler.assert_called_once()
call_args = mock_handler.call_args[0]
assert call_args[0] == str(test_file) # file_path
assert call_args[1]["bucket"] == "my-bucket" # cfg
assert call_args[2]["access_key_id"] == "AKI..." # creds
assert result["status"] == "Completed"
@patch("app.tasks.upload_to_user_integration.log_task_progress")
@patch("app.tasks.upload_to_user_integration.SessionLocal")
def test_persists_last_used_at_on_success(self, mock_session_local, mock_log_progress, tmp_path):
"""On success, last_used_at is updated and last_error is cleared."""
test_file = tmp_path / "doc.pdf"
test_file.write_bytes(b"PDF content")
from app.models import IntegrationType
from app.tasks.upload_to_user_integration import _UPLOAD_HANDLERS
integration = _make_integration(int_id=4, int_type=IntegrationType.S3)
integration.last_error = "previous error"
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = integration
mock_handler = MagicMock(return_value={"status": "Completed"})
with patch.dict(_UPLOAD_HANDLERS, {IntegrationType.S3: mock_handler}):
result = _run_upload_task(str(test_file), integration_id=4, file_id=10)
assert result["status"] == "Completed"
assert integration.last_used_at is not None
assert integration.last_error is None
@patch("app.tasks.upload_to_user_integration.log_task_progress")
@patch("app.tasks.upload_to_user_integration.SessionLocal")
def test_persists_error_and_reraises_on_failure(self, mock_session_local, mock_log_progress, tmp_path):
"""On failure, last_error is persisted on the integration and the exception is re-raised."""
test_file = tmp_path / "doc.pdf"
test_file.write_bytes(b"PDF content")
from app.models import IntegrationType
from app.tasks.upload_to_user_integration import _UPLOAD_HANDLERS
integration = _make_integration(int_id=5, int_type=IntegrationType.DROPBOX)
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = integration
mock_handler = MagicMock(side_effect=RuntimeError("Dropbox token expired"))
with patch.dict(_UPLOAD_HANDLERS, {IntegrationType.DROPBOX: mock_handler}):
with pytest.raises(RuntimeError, match="Dropbox token expired"):
_run_upload_task(str(test_file), integration_id=5, file_id=20)
assert integration.last_error == "Dropbox token expired"
@patch("app.tasks.upload_to_user_integration.log_task_progress")
@patch("app.tasks.upload_to_user_integration.SessionLocal")
def test_invalid_config_json_raises_value_error(self, mock_session_local, mock_log_progress, tmp_path):
"""ValueError is raised when integration.config contains invalid JSON."""
test_file = tmp_path / "doc.pdf"
test_file.write_bytes(b"PDF content")
from app.models import IntegrationType
integration = _make_integration(int_id=6, int_type=IntegrationType.S3)
integration.config = "NOT JSON" # corrupt config
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = integration
with pytest.raises(ValueError, match="invalid JSON in config"):
_run_upload_task(str(test_file), integration_id=6)
@patch("app.tasks.upload_to_user_integration.log_task_progress")
@patch("app.tasks.upload_to_user_integration.decrypt_value")
@patch("app.tasks.upload_to_user_integration.SessionLocal")
def test_uses_decrypt_value_for_credentials(self, mock_session_local, mock_decrypt, mock_log_progress, tmp_path):
"""credentials are decrypted using decrypt_value before being parsed as JSON."""
test_file = tmp_path / "doc.pdf"
test_file.write_bytes(b"PDF content")
from app.models import IntegrationType
from app.tasks.upload_to_user_integration import _UPLOAD_HANDLERS
integration = _make_integration(
int_id=7,
int_type=IntegrationType.S3,
config_dict={"bucket": "b", "region": "eu-west-1"},
)
# Simulate encrypted credentials stored in DB
integration.credentials = "enc:encrypted-value"
# decrypt_value should return plain JSON
mock_decrypt.return_value = json.dumps({"access_key_id": "AKI...", "secret_access_key": "S"})
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.first.return_value = integration
mock_handler = MagicMock(return_value={"status": "Completed"})
with patch.dict(_UPLOAD_HANDLERS, {IntegrationType.S3: mock_handler}):
_run_upload_task(str(test_file), integration_id=7)
mock_decrypt.assert_called_once_with("enc:encrypted-value")
# Handler should receive decrypted credentials
_, _, creds, _ = mock_handler.call_args[0]
assert creds["access_key_id"] == "AKI..."
# ---------------------------------------------------------------------------
# Tests for send_to_user_destinations task
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestSendToUserDestinations:
"""Unit tests for the send_to_user_destinations Celery task."""
def _run_task(self, file_path: str, owner_id: str, file_id: int | None = None):
"""Call the task's __wrapped__ function directly."""
from app.tasks.send_to_all import send_to_user_destinations
send_to_user_destinations.request.id = "test-task-id"
return send_to_user_destinations.__wrapped__(
file_path=file_path,
owner_id=owner_id,
file_id=file_id,
)
@patch("app.tasks.send_to_all.log_task_progress")
@patch("app.tasks.send_to_all.SessionLocal")
def test_raises_file_not_found(self, mock_session_local, mock_log_progress):
"""FileNotFoundError is raised when file does not exist."""
with pytest.raises(FileNotFoundError):
self._run_task("/nonexistent/file.pdf", owner_id="user@example.com")
@patch("app.tasks.send_to_all.log_task_progress")
@patch("app.tasks.send_to_all.SessionLocal")
def test_returns_zero_when_no_integrations(self, mock_session_local, mock_log_progress, tmp_path):
"""Returns queued=0 when there are no active DESTINATION integrations."""
test_file = tmp_path / "doc.pdf"
test_file.write_bytes(b"PDF")
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.all.return_value = []
result = self._run_task(str(test_file), owner_id="nobody@example.com", file_id=1)
assert result["queued"] == 0
assert result["status"] == "Queued"
@patch("app.tasks.send_to_all.log_task_progress")
@patch("app.tasks.send_to_all.SessionLocal")
def test_dispatches_one_task_per_integration(self, mock_session_local, mock_log_progress, tmp_path):
"""One upload_to_user_integration.delay call is made per active DESTINATION integration."""
test_file = tmp_path / "doc.pdf"
test_file.write_bytes(b"PDF")
from app.models import IntegrationType
int1 = _make_integration(int_id=10, int_type=IntegrationType.S3, name="S3 Backup")
int2 = _make_integration(int_id=11, int_type=IntegrationType.DROPBOX, name="Dropbox")
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.all.return_value = [int1, int2]
mock_celery_task = MagicMock()
mock_celery_task.delay.return_value = MagicMock(id="celery-task-id")
# The lazy import inside send_to_user_destinations uses:
# "from app.tasks.upload_to_user_integration import upload_to_user_integration"
# We must patch at the source module so the local import picks up the mock.
with patch(
"app.tasks.upload_to_user_integration.upload_to_user_integration",
mock_celery_task,
):
result = self._run_task(str(test_file), owner_id="user@example.com", file_id=99)
assert result["queued"] == 2
assert mock_celery_task.delay.call_count == 2
# Verify correct arguments
mock_celery_task.delay.assert_any_call(str(test_file), 10, 99)
mock_celery_task.delay.assert_any_call(str(test_file), 11, 99)
@patch("app.tasks.send_to_all.log_task_progress")
@patch("app.tasks.send_to_all.SessionLocal")
def test_continues_on_individual_dispatch_failure(self, mock_session_local, mock_log_progress, tmp_path):
"""If queuing one integration fails, the others are still queued."""
test_file = tmp_path / "doc.pdf"
test_file.write_bytes(b"PDF")
from app.models import IntegrationType
int1 = _make_integration(int_id=20, int_type=IntegrationType.S3, name="S3")
int2 = _make_integration(int_id=21, int_type=IntegrationType.DROPBOX, name="Dropbox")
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.all.return_value = [int1, int2]
mock_celery_task = MagicMock()
# First call fails, second succeeds
mock_celery_task.delay.side_effect = [RuntimeError("connection refused"), MagicMock(id="ok")]
with patch(
"app.tasks.upload_to_user_integration.upload_to_user_integration",
mock_celery_task,
):
result = self._run_task(str(test_file), owner_id="user@example.com", file_id=50)
# Only 1 successfully queued (the second one)
assert result["queued"] == 1
assert "integration_20_error" in result["tasks"]
# ---------------------------------------------------------------------------
# Tests for get_user_destination_count helper
# ---------------------------------------------------------------------------
@pytest.mark.unit
class TestGetUserDestinationCount:
"""Unit tests for the get_user_destination_count helper function."""
@patch("app.tasks.send_to_all.SessionLocal")
def test_returns_count_from_db(self, mock_session_local):
"""Returns the number of active DESTINATION integrations for an owner."""
from app.tasks.send_to_all import get_user_destination_count
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.count.return_value = 3
assert get_user_destination_count("user@example.com") == 3
@patch("app.tasks.send_to_all.SessionLocal")
def test_returns_zero_when_no_integrations(self, mock_session_local):
"""Returns 0 when no active DESTINATION integrations are configured."""
from app.tasks.send_to_all import get_user_destination_count
mock_db = MagicMock()
mock_session_local.return_value.__enter__.return_value = mock_db
mock_db.query.return_value.filter.return_value.count.return_value = 0
assert get_user_destination_count("empty@example.com") == 0