feat: complete issue fixes and implement remaining features
## Critical Issues (CRIT-001~003) - All Fixed
- JWT secret key validation with pydantic field_validator
- Login audit logging for success/failure attempts
- Frontend API path prefix removal
## High Priority Issues (HIGH-001~008) - All Fixed
- Project soft delete using is_active flag
- Redis session token bytes handling
- Rate limiting with slowapi (5 req/min for login)
- Attachment API permission checks
- Kanban view with drag-and-drop
- Workload heatmap UI (WorkloadPage, WorkloadHeatmap)
- TaskDetailModal integrating Comments/Attachments
- UserSelect component for task assignment
## Medium Priority Issues (MED-001~012) - All Fixed
- MED-001~005: DB commits, N+1 queries, datetime, error format, blocker flag
- MED-006: Project health dashboard (HealthService, ProjectHealthPage)
- MED-007: Capacity update API (PUT /api/users/{id}/capacity)
- MED-008: Schedule triggers (cron parsing, deadline reminders)
- MED-009: Watermark feature (image/PDF watermarking)
- MED-010~012: useEffect deps, DOM operations, PDF export
## New Files
- backend/app/api/health/ - Project health API
- backend/app/services/health_service.py
- backend/app/services/trigger_scheduler.py
- backend/app/services/watermark_service.py
- backend/app/core/rate_limiter.py
- frontend/src/pages/ProjectHealthPage.tsx
- frontend/src/components/ProjectHealthCard.tsx
- frontend/src/components/KanbanBoard.tsx
- frontend/src/components/WorkloadHeatmap.tsx
## Tests
- 113 new tests passing (health: 32, users: 14, triggers: 35, watermark: 32)
## OpenSpec Archives
- add-project-health-dashboard
- add-capacity-update-api
- add-schedule-triggers
- add-watermark-feature
- add-rate-limiting
- enhance-frontend-ux
- add-resource-management-ui
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
beabigegg
755
backend/tests/test_watermark.py
Normal file
755
backend/tests/test_watermark.py
Normal file
@@ -0,0 +1,755 @@
|
||||
"""
|
||||
Tests for MED-009: Dynamic Watermark for Downloads
|
||||
|
||||
This module contains unit tests for WatermarkService and
|
||||
integration tests for the download endpoint with watermark functionality.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import uuid
|
||||
import os
|
||||
import io
|
||||
import tempfile
|
||||
import shutil
|
||||
from datetime import datetime
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
|
||||
from app.models import User, Task, Project, Space, Attachment, AttachmentVersion
|
||||
from app.services.watermark_service import WatermarkService, watermark_service
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Test Fixtures
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_user(db):
|
||||
"""Create a test user for watermark tests."""
|
||||
user = User(
|
||||
id=str(uuid.uuid4()),
|
||||
email="watermark.test@example.com",
|
||||
employee_id="EMP-WM001",
|
||||
name="Watermark Tester",
|
||||
role_id="00000000-0000-0000-0000-000000000003",
|
||||
is_active=True,
|
||||
is_system_admin=False,
|
||||
)
|
||||
db.add(user)
|
||||
db.commit()
|
||||
return user
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_user_token(client, mock_redis, test_user):
|
||||
"""Get a token for test user."""
|
||||
from app.core.security import create_access_token, create_token_payload
|
||||
|
||||
token_data = create_token_payload(
|
||||
user_id=test_user.id,
|
||||
email=test_user.email,
|
||||
role="engineer",
|
||||
department_id=None,
|
||||
is_system_admin=False,
|
||||
)
|
||||
token = create_access_token(token_data)
|
||||
mock_redis.setex(f"session:{test_user.id}", 900, token)
|
||||
return token
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_space(db, test_user):
|
||||
"""Create a test space."""
|
||||
space = Space(
|
||||
id=str(uuid.uuid4()),
|
||||
name="Watermark Test Space",
|
||||
description="Test space for watermark tests",
|
||||
owner_id=test_user.id,
|
||||
)
|
||||
db.add(space)
|
||||
db.commit()
|
||||
return space
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_project(db, test_space, test_user):
|
||||
"""Create a test project."""
|
||||
project = Project(
|
||||
id=str(uuid.uuid4()),
|
||||
space_id=test_space.id,
|
||||
title="Watermark Test Project",
|
||||
description="Test project for watermark tests",
|
||||
owner_id=test_user.id,
|
||||
)
|
||||
db.add(project)
|
||||
db.commit()
|
||||
return project
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_task(db, test_project, test_user):
|
||||
"""Create a test task."""
|
||||
task = Task(
|
||||
id=str(uuid.uuid4()),
|
||||
project_id=test_project.id,
|
||||
title="Watermark Test Task",
|
||||
description="Test task for watermark tests",
|
||||
created_by=test_user.id,
|
||||
)
|
||||
db.add(task)
|
||||
db.commit()
|
||||
return task
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_upload_dir():
|
||||
"""Create a temporary upload directory."""
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
yield temp_dir
|
||||
shutil.rmtree(temp_dir)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_png_bytes():
|
||||
"""Create a sample PNG image as bytes."""
|
||||
img = Image.new("RGB", (200, 200), color=(255, 255, 255))
|
||||
output = io.BytesIO()
|
||||
img.save(output, format="PNG")
|
||||
output.seek(0)
|
||||
return output.getvalue()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_jpeg_bytes():
|
||||
"""Create a sample JPEG image as bytes."""
|
||||
img = Image.new("RGB", (200, 200), color=(255, 255, 255))
|
||||
output = io.BytesIO()
|
||||
img.save(output, format="JPEG")
|
||||
output.seek(0)
|
||||
return output.getvalue()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_pdf_bytes():
|
||||
"""Create a sample PDF as bytes."""
|
||||
from reportlab.lib.pagesizes import letter
|
||||
from reportlab.pdfgen import canvas
|
||||
|
||||
buffer = io.BytesIO()
|
||||
c = canvas.Canvas(buffer, pagesize=letter)
|
||||
c.drawString(100, 750, "Test PDF Document")
|
||||
c.drawString(100, 700, "This is a test page for watermarking.")
|
||||
c.showPage()
|
||||
c.drawString(100, 750, "Page 2")
|
||||
c.drawString(100, 700, "Second page content.")
|
||||
c.showPage()
|
||||
c.save()
|
||||
buffer.seek(0)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Unit Tests for WatermarkService
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestWatermarkServiceUnit:
|
||||
"""Unit tests for WatermarkService class."""
|
||||
|
||||
def test_format_watermark_text(self):
|
||||
"""Test watermark text formatting with employee_id."""
|
||||
test_time = datetime(2024, 1, 15, 10, 30, 45)
|
||||
text = WatermarkService._format_watermark_text(
|
||||
user_name="John Doe",
|
||||
employee_id="EMP001",
|
||||
download_time=test_time
|
||||
)
|
||||
|
||||
assert "John Doe" in text
|
||||
assert "EMP001" in text
|
||||
assert "2024-01-15 10:30:45" in text
|
||||
assert text == "John Doe (EMP001) - 2024-01-15 10:30:45"
|
||||
|
||||
def test_format_watermark_text_without_employee_id(self):
|
||||
"""Test that watermark text uses N/A when employee_id is not provided."""
|
||||
test_time = datetime(2024, 1, 15, 10, 30, 45)
|
||||
text = WatermarkService._format_watermark_text(
|
||||
user_name="Jane Doe",
|
||||
employee_id=None,
|
||||
download_time=test_time
|
||||
)
|
||||
|
||||
assert "Jane Doe" in text
|
||||
assert "(N/A)" in text
|
||||
assert text == "Jane Doe (N/A) - 2024-01-15 10:30:45"
|
||||
|
||||
def test_format_watermark_text_defaults_to_now(self):
|
||||
"""Test that watermark text defaults to current time."""
|
||||
text = WatermarkService._format_watermark_text(
|
||||
user_name="Jane Doe",
|
||||
employee_id="EMP002"
|
||||
)
|
||||
|
||||
assert "Jane Doe" in text
|
||||
assert "EMP002" in text
|
||||
# Should contain a date-like string
|
||||
assert "-" in text # Date separator
|
||||
|
||||
def test_is_supported_image_png(self):
|
||||
"""Test PNG is recognized as supported image."""
|
||||
service = WatermarkService()
|
||||
assert service.is_supported_image("image/png") is True
|
||||
assert service.is_supported_image("IMAGE/PNG") is True
|
||||
|
||||
def test_is_supported_image_jpeg(self):
|
||||
"""Test JPEG is recognized as supported image."""
|
||||
service = WatermarkService()
|
||||
assert service.is_supported_image("image/jpeg") is True
|
||||
assert service.is_supported_image("image/jpg") is True
|
||||
|
||||
def test_is_supported_image_unsupported(self):
|
||||
"""Test unsupported image formats are rejected."""
|
||||
service = WatermarkService()
|
||||
assert service.is_supported_image("image/gif") is False
|
||||
assert service.is_supported_image("image/bmp") is False
|
||||
assert service.is_supported_image("image/webp") is False
|
||||
|
||||
def test_is_supported_pdf(self):
|
||||
"""Test PDF is recognized."""
|
||||
service = WatermarkService()
|
||||
assert service.is_supported_pdf("application/pdf") is True
|
||||
assert service.is_supported_pdf("APPLICATION/PDF") is True
|
||||
|
||||
def test_is_supported_pdf_negative(self):
|
||||
"""Test non-PDF types are not recognized as PDF."""
|
||||
service = WatermarkService()
|
||||
assert service.is_supported_pdf("application/json") is False
|
||||
assert service.is_supported_pdf("text/plain") is False
|
||||
|
||||
def test_supports_watermark_images(self):
|
||||
"""Test supports_watermark for images."""
|
||||
service = WatermarkService()
|
||||
assert service.supports_watermark("image/png") is True
|
||||
assert service.supports_watermark("image/jpeg") is True
|
||||
|
||||
def test_supports_watermark_pdf(self):
|
||||
"""Test supports_watermark for PDF."""
|
||||
service = WatermarkService()
|
||||
assert service.supports_watermark("application/pdf") is True
|
||||
|
||||
def test_supports_watermark_unsupported(self):
|
||||
"""Test supports_watermark for unsupported types."""
|
||||
service = WatermarkService()
|
||||
assert service.supports_watermark("text/plain") is False
|
||||
assert service.supports_watermark("application/zip") is False
|
||||
assert service.supports_watermark("application/octet-stream") is False
|
||||
|
||||
|
||||
class TestImageWatermarking:
|
||||
"""Unit tests for image watermarking functionality."""
|
||||
|
||||
def test_add_image_watermark_png(self, sample_png_bytes):
|
||||
"""Test adding watermark to PNG image."""
|
||||
test_time = datetime(2024, 1, 15, 10, 30, 45)
|
||||
|
||||
result_bytes, output_format = watermark_service.add_image_watermark(
|
||||
image_bytes=sample_png_bytes,
|
||||
user_name="Test User",
|
||||
employee_id="EMP001",
|
||||
download_time=test_time
|
||||
)
|
||||
|
||||
# Verify output is valid image bytes
|
||||
assert len(result_bytes) > 0
|
||||
assert output_format.lower() == "png"
|
||||
|
||||
# Verify output is valid PNG image
|
||||
result_image = Image.open(io.BytesIO(result_bytes))
|
||||
assert result_image.format == "PNG"
|
||||
assert result_image.size == (200, 200)
|
||||
|
||||
def test_add_image_watermark_jpeg(self, sample_jpeg_bytes):
|
||||
"""Test adding watermark to JPEG image."""
|
||||
test_time = datetime(2024, 1, 15, 10, 30, 45)
|
||||
|
||||
result_bytes, output_format = watermark_service.add_image_watermark(
|
||||
image_bytes=sample_jpeg_bytes,
|
||||
user_name="Test User",
|
||||
employee_id="EMP001",
|
||||
download_time=test_time
|
||||
)
|
||||
|
||||
# Verify output is valid image bytes
|
||||
assert len(result_bytes) > 0
|
||||
assert output_format.lower() == "jpeg"
|
||||
|
||||
# Verify output is valid JPEG image
|
||||
result_image = Image.open(io.BytesIO(result_bytes))
|
||||
assert result_image.format == "JPEG"
|
||||
assert result_image.size == (200, 200)
|
||||
|
||||
def test_add_image_watermark_preserves_dimensions(self, sample_png_bytes):
|
||||
"""Test that watermarking preserves image dimensions."""
|
||||
original = Image.open(io.BytesIO(sample_png_bytes))
|
||||
original_size = original.size
|
||||
|
||||
result_bytes, _ = watermark_service.add_image_watermark(
|
||||
image_bytes=sample_png_bytes,
|
||||
user_name="Test User",
|
||||
employee_id="EMP001"
|
||||
)
|
||||
|
||||
result = Image.open(io.BytesIO(result_bytes))
|
||||
assert result.size == original_size
|
||||
|
||||
def test_add_image_watermark_modifies_image(self, sample_png_bytes):
|
||||
"""Test that watermark actually modifies the image."""
|
||||
result_bytes, _ = watermark_service.add_image_watermark(
|
||||
image_bytes=sample_png_bytes,
|
||||
user_name="Test User",
|
||||
employee_id="EMP001"
|
||||
)
|
||||
|
||||
# The watermarked image should be different from original
|
||||
# (Note: size might differ slightly due to compression)
|
||||
# We verify the image data is actually different
|
||||
original = Image.open(io.BytesIO(sample_png_bytes))
|
||||
result = Image.open(io.BytesIO(result_bytes))
|
||||
|
||||
# Convert to same mode for comparison
|
||||
original_rgb = original.convert("RGB")
|
||||
result_rgb = result.convert("RGB")
|
||||
|
||||
# Compare pixel data - they should be different
|
||||
original_data = list(original_rgb.getdata())
|
||||
result_data = list(result_rgb.getdata())
|
||||
|
||||
# At least some pixels should be different (watermark added)
|
||||
different_pixels = sum(1 for o, r in zip(original_data, result_data) if o != r)
|
||||
assert different_pixels > 0, "Watermark should modify image pixels"
|
||||
|
||||
def test_add_image_watermark_large_image(self):
|
||||
"""Test watermarking a larger image."""
|
||||
# Create a larger image
|
||||
large_img = Image.new("RGB", (1920, 1080), color=(100, 150, 200))
|
||||
output = io.BytesIO()
|
||||
large_img.save(output, format="PNG")
|
||||
large_bytes = output.getvalue()
|
||||
|
||||
result_bytes, output_format = watermark_service.add_image_watermark(
|
||||
image_bytes=large_bytes,
|
||||
user_name="Large Image User",
|
||||
employee_id="EMP-LARGE"
|
||||
)
|
||||
|
||||
assert len(result_bytes) > 0
|
||||
result_image = Image.open(io.BytesIO(result_bytes))
|
||||
assert result_image.size == (1920, 1080)
|
||||
|
||||
|
||||
class TestPdfWatermarking:
|
||||
"""Unit tests for PDF watermarking functionality."""
|
||||
|
||||
def test_add_pdf_watermark_basic(self, sample_pdf_bytes):
|
||||
"""Test adding watermark to PDF."""
|
||||
import fitz # PyMuPDF
|
||||
|
||||
test_time = datetime(2024, 1, 15, 10, 30, 45)
|
||||
|
||||
result_bytes = watermark_service.add_pdf_watermark(
|
||||
pdf_bytes=sample_pdf_bytes,
|
||||
user_name="PDF Test User",
|
||||
employee_id="EMP-PDF001",
|
||||
download_time=test_time
|
||||
)
|
||||
|
||||
# Verify output is valid PDF bytes
|
||||
assert len(result_bytes) > 0
|
||||
|
||||
# Verify output is valid PDF using PyMuPDF
|
||||
result_pdf = fitz.open(stream=result_bytes, filetype="pdf")
|
||||
assert len(result_pdf) == 2
|
||||
result_pdf.close()
|
||||
|
||||
def test_add_pdf_watermark_preserves_page_count(self, sample_pdf_bytes):
|
||||
"""Test that watermarking preserves page count."""
|
||||
import fitz # PyMuPDF
|
||||
|
||||
original_pdf = fitz.open(stream=sample_pdf_bytes, filetype="pdf")
|
||||
original_page_count = len(original_pdf)
|
||||
original_pdf.close()
|
||||
|
||||
result_bytes = watermark_service.add_pdf_watermark(
|
||||
pdf_bytes=sample_pdf_bytes,
|
||||
user_name="Test User",
|
||||
employee_id="EMP001"
|
||||
)
|
||||
|
||||
result_pdf = fitz.open(stream=result_bytes, filetype="pdf")
|
||||
assert len(result_pdf) == original_page_count
|
||||
result_pdf.close()
|
||||
|
||||
def test_add_pdf_watermark_modifies_content(self, sample_pdf_bytes):
|
||||
"""Test that watermark actually modifies the PDF content."""
|
||||
result_bytes = watermark_service.add_pdf_watermark(
|
||||
pdf_bytes=sample_pdf_bytes,
|
||||
user_name="Modified User",
|
||||
employee_id="EMP-MOD"
|
||||
)
|
||||
|
||||
# The watermarked PDF should be different from original
|
||||
assert result_bytes != sample_pdf_bytes
|
||||
|
||||
def test_add_pdf_watermark_single_page(self):
|
||||
"""Test watermarking a single-page PDF."""
|
||||
import fitz # PyMuPDF
|
||||
|
||||
# Create single page PDF with PyMuPDF
|
||||
doc = fitz.open()
|
||||
page = doc.new_page(width=612, height=792) # Letter size
|
||||
page.insert_text(point=(100, 750), text="Single Page Document", fontsize=12)
|
||||
buffer = io.BytesIO()
|
||||
doc.save(buffer)
|
||||
doc.close()
|
||||
single_page_bytes = buffer.getvalue()
|
||||
|
||||
result_bytes = watermark_service.add_pdf_watermark(
|
||||
pdf_bytes=single_page_bytes,
|
||||
user_name="Single Page User",
|
||||
employee_id="EMP-SINGLE"
|
||||
)
|
||||
|
||||
result_pdf = fitz.open(stream=result_bytes, filetype="pdf")
|
||||
assert len(result_pdf) == 1
|
||||
result_pdf.close()
|
||||
|
||||
def test_add_pdf_watermark_many_pages(self):
|
||||
"""Test watermarking a multi-page PDF."""
|
||||
import fitz # PyMuPDF
|
||||
|
||||
# Create multi-page PDF with PyMuPDF
|
||||
doc = fitz.open()
|
||||
for i in range(5):
|
||||
page = doc.new_page(width=612, height=792)
|
||||
page.insert_text(point=(100, 750), text=f"Page {i + 1}", fontsize=12)
|
||||
buffer = io.BytesIO()
|
||||
doc.save(buffer)
|
||||
doc.close()
|
||||
multi_page_bytes = buffer.getvalue()
|
||||
|
||||
result_bytes = watermark_service.add_pdf_watermark(
|
||||
pdf_bytes=multi_page_bytes,
|
||||
user_name="Multi Page User",
|
||||
employee_id="EMP-MULTI"
|
||||
)
|
||||
|
||||
result_pdf = fitz.open(stream=result_bytes, filetype="pdf")
|
||||
assert len(result_pdf) == 5
|
||||
result_pdf.close()
|
||||
|
||||
|
||||
class TestWatermarkServiceConfiguration:
|
||||
"""Tests for WatermarkService configuration constants."""
|
||||
|
||||
def test_default_opacity(self):
|
||||
"""Test default watermark opacity."""
|
||||
assert WatermarkService.WATERMARK_OPACITY == 0.3
|
||||
|
||||
def test_default_angle(self):
|
||||
"""Test default watermark angle."""
|
||||
assert WatermarkService.WATERMARK_ANGLE == -45
|
||||
|
||||
def test_default_font_size(self):
|
||||
"""Test default watermark font size."""
|
||||
assert WatermarkService.WATERMARK_FONT_SIZE == 24
|
||||
|
||||
def test_default_color(self):
|
||||
"""Test default watermark color (gray)."""
|
||||
assert WatermarkService.WATERMARK_COLOR == (128, 128, 128)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Integration Tests for Download with Watermark
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestDownloadWithWatermark:
|
||||
"""Integration tests for download endpoint with watermark."""
|
||||
|
||||
def test_download_png_with_watermark(
|
||||
self, client, test_user_token, test_task, db, monkeypatch, temp_upload_dir, sample_png_bytes
|
||||
):
|
||||
"""Test downloading PNG file applies watermark."""
|
||||
from pathlib import Path
|
||||
from app.services.file_storage_service import file_storage_service
|
||||
monkeypatch.setattr("app.core.config.settings.UPLOAD_DIR", temp_upload_dir)
|
||||
monkeypatch.setattr(file_storage_service, "base_dir", Path(temp_upload_dir))
|
||||
|
||||
# Create attachment and version
|
||||
attachment_id = str(uuid.uuid4())
|
||||
version_id = str(uuid.uuid4())
|
||||
|
||||
# Save the file to disk
|
||||
file_dir = os.path.join(temp_upload_dir, test_task.project_id, test_task.id, attachment_id, "v1")
|
||||
os.makedirs(file_dir, exist_ok=True)
|
||||
file_path = os.path.join(file_dir, "test.png")
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(sample_png_bytes)
|
||||
|
||||
relative_path = os.path.join(test_task.project_id, test_task.id, attachment_id, "v1", "test.png")
|
||||
|
||||
attachment = Attachment(
|
||||
id=attachment_id,
|
||||
task_id=test_task.id,
|
||||
filename="test.png",
|
||||
original_filename="test.png",
|
||||
mime_type="image/png",
|
||||
file_size=len(sample_png_bytes),
|
||||
current_version=1,
|
||||
uploaded_by=test_task.created_by,
|
||||
)
|
||||
db.add(attachment)
|
||||
|
||||
version = AttachmentVersion(
|
||||
id=version_id,
|
||||
attachment_id=attachment_id,
|
||||
version=1,
|
||||
file_path=relative_path,
|
||||
file_size=len(sample_png_bytes),
|
||||
checksum="0" * 64,
|
||||
uploaded_by=test_task.created_by,
|
||||
)
|
||||
db.add(version)
|
||||
db.commit()
|
||||
|
||||
# Download the file
|
||||
response = client.get(
|
||||
f"/api/attachments/{attachment_id}/download",
|
||||
headers={"Authorization": f"Bearer {test_user_token}"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "image/png"
|
||||
|
||||
# Verify watermark was applied (image should be different)
|
||||
downloaded_image = Image.open(io.BytesIO(response.content))
|
||||
original_image = Image.open(io.BytesIO(sample_png_bytes))
|
||||
|
||||
# Convert to comparable format
|
||||
downloaded_rgb = downloaded_image.convert("RGB")
|
||||
original_rgb = original_image.convert("RGB")
|
||||
|
||||
downloaded_data = list(downloaded_rgb.getdata())
|
||||
original_data = list(original_rgb.getdata())
|
||||
|
||||
# At least some pixels should be different (watermark present)
|
||||
different_pixels = sum(1 for o, d in zip(original_data, downloaded_data) if o != d)
|
||||
assert different_pixels > 0, "Downloaded image should have watermark"
|
||||
|
||||
def test_download_pdf_with_watermark(
|
||||
self, client, test_user_token, test_task, db, monkeypatch, temp_upload_dir, sample_pdf_bytes
|
||||
):
|
||||
"""Test downloading PDF file applies watermark."""
|
||||
from pathlib import Path
|
||||
from app.services.file_storage_service import file_storage_service
|
||||
monkeypatch.setattr("app.core.config.settings.UPLOAD_DIR", temp_upload_dir)
|
||||
monkeypatch.setattr(file_storage_service, "base_dir", Path(temp_upload_dir))
|
||||
|
||||
# Create attachment and version
|
||||
attachment_id = str(uuid.uuid4())
|
||||
version_id = str(uuid.uuid4())
|
||||
|
||||
# Save the file to disk
|
||||
file_dir = os.path.join(temp_upload_dir, test_task.project_id, test_task.id, attachment_id, "v1")
|
||||
os.makedirs(file_dir, exist_ok=True)
|
||||
file_path = os.path.join(file_dir, "test.pdf")
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(sample_pdf_bytes)
|
||||
|
||||
relative_path = os.path.join(test_task.project_id, test_task.id, attachment_id, "v1", "test.pdf")
|
||||
|
||||
attachment = Attachment(
|
||||
id=attachment_id,
|
||||
task_id=test_task.id,
|
||||
filename="test.pdf",
|
||||
original_filename="test.pdf",
|
||||
mime_type="application/pdf",
|
||||
file_size=len(sample_pdf_bytes),
|
||||
current_version=1,
|
||||
uploaded_by=test_task.created_by,
|
||||
)
|
||||
db.add(attachment)
|
||||
|
||||
version = AttachmentVersion(
|
||||
id=version_id,
|
||||
attachment_id=attachment_id,
|
||||
version=1,
|
||||
file_path=relative_path,
|
||||
file_size=len(sample_pdf_bytes),
|
||||
checksum="0" * 64,
|
||||
uploaded_by=test_task.created_by,
|
||||
)
|
||||
db.add(version)
|
||||
db.commit()
|
||||
|
||||
# Download the file
|
||||
response = client.get(
|
||||
f"/api/attachments/{attachment_id}/download",
|
||||
headers={"Authorization": f"Bearer {test_user_token}"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "application/pdf"
|
||||
|
||||
# Verify watermark was applied (PDF content should be different)
|
||||
assert response.content != sample_pdf_bytes, "Downloaded PDF should have watermark"
|
||||
|
||||
def test_download_unsupported_file_no_watermark(
|
||||
self, client, test_user_token, test_task, db, monkeypatch, temp_upload_dir
|
||||
):
|
||||
"""Test downloading unsupported file type returns original without watermark."""
|
||||
from pathlib import Path
|
||||
from app.services.file_storage_service import file_storage_service
|
||||
monkeypatch.setattr("app.core.config.settings.UPLOAD_DIR", temp_upload_dir)
|
||||
monkeypatch.setattr(file_storage_service, "base_dir", Path(temp_upload_dir))
|
||||
|
||||
# Create a text file
|
||||
text_content = b"This is a plain text file."
|
||||
|
||||
attachment_id = str(uuid.uuid4())
|
||||
version_id = str(uuid.uuid4())
|
||||
|
||||
# Save the file to disk
|
||||
file_dir = os.path.join(temp_upload_dir, test_task.project_id, test_task.id, attachment_id, "v1")
|
||||
os.makedirs(file_dir, exist_ok=True)
|
||||
file_path = os.path.join(file_dir, "test.txt")
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(text_content)
|
||||
|
||||
relative_path = os.path.join(test_task.project_id, test_task.id, attachment_id, "v1", "test.txt")
|
||||
|
||||
attachment = Attachment(
|
||||
id=attachment_id,
|
||||
task_id=test_task.id,
|
||||
filename="test.txt",
|
||||
original_filename="test.txt",
|
||||
mime_type="text/plain",
|
||||
file_size=len(text_content),
|
||||
current_version=1,
|
||||
uploaded_by=test_task.created_by,
|
||||
)
|
||||
db.add(attachment)
|
||||
|
||||
version = AttachmentVersion(
|
||||
id=version_id,
|
||||
attachment_id=attachment_id,
|
||||
version=1,
|
||||
file_path=relative_path,
|
||||
file_size=len(text_content),
|
||||
checksum="0" * 64,
|
||||
uploaded_by=test_task.created_by,
|
||||
)
|
||||
db.add(version)
|
||||
db.commit()
|
||||
|
||||
# Download the file
|
||||
response = client.get(
|
||||
f"/api/attachments/{attachment_id}/download",
|
||||
headers={"Authorization": f"Bearer {test_user_token}"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
# Content should be unchanged for unsupported types
|
||||
assert response.content == text_content
|
||||
|
||||
def test_download_jpeg_with_watermark(
|
||||
self, client, test_user_token, test_task, db, monkeypatch, temp_upload_dir, sample_jpeg_bytes
|
||||
):
|
||||
"""Test downloading JPEG file applies watermark."""
|
||||
from pathlib import Path
|
||||
from app.services.file_storage_service import file_storage_service
|
||||
monkeypatch.setattr("app.core.config.settings.UPLOAD_DIR", temp_upload_dir)
|
||||
monkeypatch.setattr(file_storage_service, "base_dir", Path(temp_upload_dir))
|
||||
|
||||
attachment_id = str(uuid.uuid4())
|
||||
version_id = str(uuid.uuid4())
|
||||
|
||||
# Save the file to disk
|
||||
file_dir = os.path.join(temp_upload_dir, test_task.project_id, test_task.id, attachment_id, "v1")
|
||||
os.makedirs(file_dir, exist_ok=True)
|
||||
file_path = os.path.join(file_dir, "test.jpg")
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(sample_jpeg_bytes)
|
||||
|
||||
relative_path = os.path.join(test_task.project_id, test_task.id, attachment_id, "v1", "test.jpg")
|
||||
|
||||
attachment = Attachment(
|
||||
id=attachment_id,
|
||||
task_id=test_task.id,
|
||||
filename="test.jpg",
|
||||
original_filename="test.jpg",
|
||||
mime_type="image/jpeg",
|
||||
file_size=len(sample_jpeg_bytes),
|
||||
current_version=1,
|
||||
uploaded_by=test_task.created_by,
|
||||
)
|
||||
db.add(attachment)
|
||||
|
||||
version = AttachmentVersion(
|
||||
id=version_id,
|
||||
attachment_id=attachment_id,
|
||||
version=1,
|
||||
file_path=relative_path,
|
||||
file_size=len(sample_jpeg_bytes),
|
||||
checksum="0" * 64,
|
||||
uploaded_by=test_task.created_by,
|
||||
)
|
||||
db.add(version)
|
||||
db.commit()
|
||||
|
||||
# Download the file
|
||||
response = client.get(
|
||||
f"/api/attachments/{attachment_id}/download",
|
||||
headers={"Authorization": f"Bearer {test_user_token}"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "image/jpeg"
|
||||
|
||||
# Verify the response is a valid JPEG
|
||||
downloaded_image = Image.open(io.BytesIO(response.content))
|
||||
assert downloaded_image.format == "JPEG"
|
||||
|
||||
|
||||
class TestWatermarkErrorHandling:
|
||||
"""Tests for watermark error handling and graceful degradation."""
|
||||
|
||||
def test_watermark_service_singleton_exists(self):
|
||||
"""Test that watermark_service singleton is available."""
|
||||
assert watermark_service is not None
|
||||
assert isinstance(watermark_service, WatermarkService)
|
||||
|
||||
def test_invalid_image_bytes_graceful_handling(self):
|
||||
"""Test handling of invalid image bytes."""
|
||||
invalid_bytes = b"not an image"
|
||||
|
||||
with pytest.raises(Exception):
|
||||
# Should raise an exception for invalid image data
|
||||
watermark_service.add_image_watermark(
|
||||
image_bytes=invalid_bytes,
|
||||
user_name="Test",
|
||||
employee_id="EMP001"
|
||||
)
|
||||
|
||||
def test_invalid_pdf_bytes_graceful_handling(self):
|
||||
"""Test handling of invalid PDF bytes."""
|
||||
invalid_bytes = b"not a pdf"
|
||||
|
||||
with pytest.raises(Exception):
|
||||
# Should raise an exception for invalid PDF data
|
||||
watermark_service.add_pdf_watermark(
|
||||
pdf_bytes=invalid_bytes,
|
||||
user_name="Test",
|
||||
employee_id="EMP001"
|
||||
)
|
||||
Reference in New Issue
Block a user