feat(tasks): dynamic routing to user-specific destination integrations

- Add app/tasks/upload_to_user_integration.py: new Celery task that
  uploads a processed document to a specific UserIntegration using its
  own per-user config and Fernet-decrypted credentials. Supports all
  DESTINATION types: Dropbox, S3, Google Drive, OneDrive, WebDAV,
  Nextcloud, FTP, SFTP, Paperless-ngx, Email (SMTP), and Rclone.

- Extend app/tasks/send_to_all.py: add send_to_user_destinations task
  (queries active DESTINATION UserIntegrations for an owner and
  dispatches one upload_to_user_integration task per integration) and
  get_user_destination_count helper used by finalize_document_storage.

- Refactor app/tasks/finalize_document_storage.py: after processing,
  look up the document owner; if the owner has active DESTINATION
  integrations route exclusively to those (user-specific routing),
  otherwise fall back to the global send_to_all_destinations.

- Update tests/test_finalize_storage.py: add autouse fixture to prevent
  Redis hangs, update all existing tests with new mock parameters, add
  TestFinalizeDocumentStorageUserRouting class with four new tests that
  validate user-specific vs global routing decisions.

- Add tests/test_user_integration_upload.py: 14 new unit tests covering
  upload_to_user_integration (handler dispatch, error persistence,
  last_used_at update, credential decryption, skip for unknown types)
  and send_to_user_destinations / get_user_destination_count.

- Update docs/StorageArchitecture.md: document the user-specific
  destination routing feature, supported types, multiple-destination
  behaviour, and global fallback semantics.

Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
2026-03-08 21:24:14 +00:00
parent 5d3fe88be9
commit 89bf07d2ef
6 changed files with 1579 additions and 60 deletions
+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()
)
+715
View File
@@ -0,0 +1,715 @@
#!/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:
with smtplib.SMTP(host, port, timeout=30) as smtp:
smtp.starttls() # nosec B608
if username and password:
smtp.login(username, password)
smtp.sendmail(msg["From"], [recipient], msg.as_string())
else:
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."""
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")
# Write the user's rclone config to a temp file so we don't touch the system config
import tempfile
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
+241 -42
View File
@@ -6,13 +6,33 @@ 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.mark.unit
class TestFinalizeDocumentStorage:
"""Tests for finalize_document_storage Celery task."""
@pytest.fixture(autouse=True)
def _patch_celery_tasks(self, mocker):
"""Prevent all lazy-imported Celery tasks from actually connecting to Redis."""
mocker.patch("app.tasks.compute_embedding.compute_document_embedding")
mocker.patch("app.tasks.convert_to_pdfa.convert_to_pdfa", create=True)
@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 +41,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 +54,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 +77,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 +90,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 +105,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 +140,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 +155,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 +184,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 +203,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 +227,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 +245,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 +257,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 +271,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 +286,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 +298,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 +309,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 +327,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 +355,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 +374,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 +398,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 +412,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 +449,191 @@ 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."""
@pytest.fixture(autouse=True)
def _patch_celery_tasks(self, mocker):
"""Prevent all lazy-imported Celery tasks from actually connecting to Redis."""
mocker.patch("app.tasks.compute_embedding.compute_document_embedding")
mocker.patch("app.tasks.convert_to_pdfa.convert_to_pdfa", create=True)
@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"
+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