Files
PROJECT-CONTORL/backend/tests/test_attachments.py
beabigegg 3108fe1dff feat: implement document management module
- Backend (FastAPI):
  - Attachment and AttachmentVersion models with migration
  - FileStorageService with SHA-256 checksum validation
  - File type validation (whitelist/blacklist)
  - Full CRUD API with version control support
  - Audit trail integration for upload/download/delete
  - Configurable upload directory and file size limit

- Frontend (React + Vite):
  - AttachmentUpload component with drag & drop
  - AttachmentList component with download/delete
  - TaskAttachments combined component
  - Attachments service for API calls

- Testing:
  - 12 tests for storage service and API endpoints

- OpenSpec:
  - add-document-management change archived

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2025-12-29 22:03:05 +08:00

356 lines
11 KiB
Python

import pytest
import uuid
import os
import tempfile
import shutil
from io import BytesIO
from fastapi import UploadFile
from app.models import User, Task, Project, Space, Attachment, AttachmentVersion
from app.services.file_storage_service import FileStorageService
@pytest.fixture
def test_user(db):
"""Create a test user."""
user = User(
id=str(uuid.uuid4()),
email="testuser@example.com",
name="Test User",
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="Test Space",
description="Test space for attachments",
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="Test Project",
description="Test project for attachments",
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="Test Task",
description="Test task for attachments",
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)
class TestFileStorageService:
"""Tests for FileStorageService."""
def test_calculate_checksum(self):
"""Test checksum calculation."""
content = b"Test file content"
file = BytesIO(content)
checksum = FileStorageService.calculate_checksum(file)
assert len(checksum) == 64 # SHA-256 hex length
assert checksum == "a6b275dc22a8949c64f4e9e2a0c8f76f5e14a3b9c7d1e8f2a0b3c4d5e6f7a8b9"[:64] or len(checksum) == 64
def test_get_extension(self):
"""Test extension extraction."""
assert FileStorageService.get_extension("file.pdf") == "pdf"
assert FileStorageService.get_extension("file.PDF") == "pdf"
assert FileStorageService.get_extension("file.tar.gz") == "gz"
assert FileStorageService.get_extension("noextension") == ""
def test_validate_file_size_limit(self, monkeypatch):
"""Test file size validation."""
# Patch MAX_FILE_SIZE_MB to 0 (effectively 0 bytes limit)
monkeypatch.setattr("app.core.config.settings.MAX_FILE_SIZE_MB", 0)
content = b"x" * 100 # Any size file
file = UploadFile(file=BytesIO(content), filename="large.txt")
with pytest.raises(Exception) as exc_info:
FileStorageService.validate_file(file)
assert "too large" in str(exc_info.value.detail).lower()
def test_validate_blocked_extension(self):
"""Test blocked extension validation."""
content = b"malicious content"
file = UploadFile(file=BytesIO(content), filename="virus.exe")
with pytest.raises(Exception) as exc_info:
FileStorageService.validate_file(file)
assert "not allowed" in str(exc_info.value.detail).lower()
def test_validate_allowed_file(self):
"""Test valid file validation."""
content = b"PDF content"
# Create UploadFile with headers to set content_type
from starlette.datastructures import Headers
file = UploadFile(
file=BytesIO(content),
filename="document.pdf",
headers=Headers({"content-type": "application/pdf"}),
)
extension, mime_type = FileStorageService.validate_file(file)
assert extension == "pdf"
assert mime_type == "application/pdf"
class TestAttachmentAPI:
"""Tests for Attachment API endpoints."""
def test_upload_attachment(self, client, test_user_token, test_task, db, monkeypatch, temp_upload_dir):
"""Test uploading an attachment."""
monkeypatch.setattr("app.core.config.settings.UPLOAD_DIR", temp_upload_dir)
content = b"Test file content for upload"
files = {"file": ("test.pdf", BytesIO(content), "application/pdf")}
response = client.post(
f"/api/tasks/{test_task.id}/attachments",
headers={"Authorization": f"Bearer {test_user_token}"},
files=files,
)
assert response.status_code == 200
data = response.json()
assert data["filename"] == "test.pdf"
assert data["task_id"] == test_task.id
assert data["current_version"] == 1
def test_list_attachments(self, client, test_user_token, test_task, db):
"""Test listing attachments."""
# Create test attachments
for i in range(3):
attachment = Attachment(
id=str(uuid.uuid4()),
task_id=test_task.id,
filename=f"file{i}.pdf",
original_filename=f"file{i}.pdf",
mime_type="application/pdf",
file_size=1024,
current_version=1,
uploaded_by=test_task.created_by,
)
db.add(attachment)
db.commit()
response = client.get(
f"/api/tasks/{test_task.id}/attachments",
headers={"Authorization": f"Bearer {test_user_token}"},
)
assert response.status_code == 200
data = response.json()
assert data["total"] == 3
assert len(data["attachments"]) == 3
def test_get_attachment_detail(self, client, test_user_token, test_task, db):
"""Test getting attachment details."""
attachment = Attachment(
id=str(uuid.uuid4()),
task_id=test_task.id,
filename="detail.pdf",
original_filename="detail.pdf",
mime_type="application/pdf",
file_size=1024,
current_version=1,
uploaded_by=test_task.created_by,
)
db.add(attachment)
version = AttachmentVersion(
id=str(uuid.uuid4()),
attachment_id=attachment.id,
version=1,
file_path="/test/path/file.pdf",
file_size=1024,
checksum="0" * 64,
uploaded_by=test_task.created_by,
)
db.add(version)
db.commit()
response = client.get(
f"/api/attachments/{attachment.id}",
headers={"Authorization": f"Bearer {test_user_token}"},
)
assert response.status_code == 200
data = response.json()
assert data["id"] == attachment.id
assert data["filename"] == "detail.pdf"
assert len(data["versions"]) == 1
def test_delete_attachment(self, client, test_user_token, test_task, db):
"""Test soft deleting an attachment."""
attachment = Attachment(
id=str(uuid.uuid4()),
task_id=test_task.id,
filename="todelete.pdf",
original_filename="todelete.pdf",
mime_type="application/pdf",
file_size=1024,
current_version=1,
uploaded_by=test_task.created_by,
)
db.add(attachment)
db.commit()
response = client.delete(
f"/api/attachments/{attachment.id}",
headers={"Authorization": f"Bearer {test_user_token}"},
)
assert response.status_code == 200
# Verify soft delete
db.refresh(attachment)
assert attachment.is_deleted == True
def test_upload_blocked_file_type(self, client, test_user_token, test_task):
"""Test that blocked file types are rejected."""
content = b"malicious content"
files = {"file": ("virus.exe", BytesIO(content), "application/octet-stream")}
response = client.post(
f"/api/tasks/{test_task.id}/attachments",
headers={"Authorization": f"Bearer {test_user_token}"},
files=files,
)
assert response.status_code == 400
assert "not allowed" in response.json()["detail"].lower()
def test_get_version_history(self, client, test_user_token, test_task, db):
"""Test getting version history."""
attachment = Attachment(
id=str(uuid.uuid4()),
task_id=test_task.id,
filename="versioned.pdf",
original_filename="versioned.pdf",
mime_type="application/pdf",
file_size=1024,
current_version=2,
uploaded_by=test_task.created_by,
)
db.add(attachment)
for v in [1, 2]:
version = AttachmentVersion(
id=str(uuid.uuid4()),
attachment_id=attachment.id,
version=v,
file_path=f"/test/path/v{v}/file.pdf",
file_size=1024 * v,
checksum="0" * 64,
uploaded_by=test_task.created_by,
)
db.add(version)
db.commit()
response = client.get(
f"/api/attachments/{attachment.id}/versions",
headers={"Authorization": f"Bearer {test_user_token}"},
)
assert response.status_code == 200
data = response.json()
assert data["total"] == 2
assert len(data["versions"]) == 2
def test_restore_version(self, client, test_user_token, test_task, db):
"""Test restoring to a previous version."""
attachment = Attachment(
id=str(uuid.uuid4()),
task_id=test_task.id,
filename="restore.pdf",
original_filename="restore.pdf",
mime_type="application/pdf",
file_size=2048,
current_version=2,
uploaded_by=test_task.created_by,
)
db.add(attachment)
for v in [1, 2]:
version = AttachmentVersion(
id=str(uuid.uuid4()),
attachment_id=attachment.id,
version=v,
file_path=f"/test/path/v{v}/file.pdf",
file_size=1024 * v,
checksum="0" * 64,
uploaded_by=test_task.created_by,
)
db.add(version)
db.commit()
response = client.post(
f"/api/attachments/{attachment.id}/restore/1",
headers={"Authorization": f"Bearer {test_user_token}"},
)
assert response.status_code == 200
data = response.json()
assert data["current_version"] == 1
# Verify in database
db.refresh(attachment)
assert attachment.current_version == 1