Code review feedback: improve readability of S3 upload call
Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
@@ -10,8 +10,6 @@ from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.celery_app import celery
|
||||
from app.utils import log_task_progress
|
||||
from app.database import SessionLocal
|
||||
from app.models import FileRecord
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -229,7 +227,13 @@ def upload_to_onedrive(self, file_path: str, file_id: int = None):
|
||||
"""
|
||||
task_id = self.request.id
|
||||
logger.info(f"[{task_id}] Starting OneDrive upload: {file_path}")
|
||||
log_task_progress(task_id, "upload_to_onedrive", "in_progress", f"Uploading to OneDrive: {os.path.basename(file_path)}", file_id=file_id)
|
||||
log_task_progress(
|
||||
task_id,
|
||||
"upload_to_onedrive",
|
||||
"in_progress",
|
||||
f"Uploading to OneDrive: {os.path.basename(file_path)}",
|
||||
file_id=file_id,
|
||||
)
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
error_msg = f"File not found: {file_path}"
|
||||
@@ -261,7 +265,9 @@ def upload_to_onedrive(self, file_path: str, file_id: int = None):
|
||||
web_url = result.get("webUrl", "Not available")
|
||||
logger.info(f"[{task_id}] Successfully uploaded {filename} to OneDrive at path {settings.onedrive_folder_path}")
|
||||
logger.info(f"[{task_id}] File accessible at: {web_url}")
|
||||
log_task_progress(task_id, "upload_to_onedrive", "success", f"Uploaded to OneDrive: {filename}", file_id=file_id)
|
||||
log_task_progress(
|
||||
task_id, "upload_to_onedrive", "success", f"Uploaded to OneDrive: {filename}", file_id=file_id
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "Completed",
|
||||
|
||||
+26
-25
@@ -8,24 +8,25 @@ from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.celery_app import celery
|
||||
from app.utils import log_task_progress
|
||||
from app.database import SessionLocal
|
||||
from app.models import FileRecord
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry, bind=True)
|
||||
def upload_to_s3(self, file_path: str, file_id: int = None):
|
||||
"""
|
||||
Uploads a file to Amazon S3 in the configured bucket and folder.
|
||||
|
||||
|
||||
Args:
|
||||
file_path: Path to the file to upload
|
||||
file_id: Optional file ID to associate with logs
|
||||
"""
|
||||
task_id = self.request.id
|
||||
logger.info(f"[{task_id}] Starting S3 upload: {file_path}")
|
||||
log_task_progress(task_id, "upload_to_s3", "in_progress", f"Uploading to S3: {os.path.basename(file_path)}", file_id=file_id)
|
||||
|
||||
log_task_progress(
|
||||
task_id, "upload_to_s3", "in_progress", f"Uploading to S3: {os.path.basename(file_path)}", file_id=file_id
|
||||
)
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
error_msg = f"File not found: {file_path}"
|
||||
logger.error(f"[{task_id}] {error_msg}")
|
||||
@@ -51,59 +52,59 @@ def upload_to_s3(self, file_path: str, file_id: int = None):
|
||||
try:
|
||||
# Create S3 client
|
||||
s3_client = boto3.client(
|
||||
's3',
|
||||
"s3",
|
||||
region_name=settings.aws_region,
|
||||
aws_access_key_id=settings.aws_access_key_id,
|
||||
aws_secret_access_key=settings.aws_secret_access_key
|
||||
aws_secret_access_key=settings.aws_secret_access_key,
|
||||
)
|
||||
|
||||
|
||||
# Construct the S3 key (path within the bucket)
|
||||
if settings.s3_folder_prefix:
|
||||
# Ensure folder prefix ends with a slash
|
||||
folder_prefix = settings.s3_folder_prefix
|
||||
if not folder_prefix.endswith('/'):
|
||||
folder_prefix += '/'
|
||||
if not folder_prefix.endswith("/"):
|
||||
folder_prefix += "/"
|
||||
s3_key = f"{folder_prefix}{filename}"
|
||||
else:
|
||||
s3_key = filename
|
||||
|
||||
|
||||
# Prepare extra arguments
|
||||
extra_args = {
|
||||
'StorageClass': settings.s3_storage_class
|
||||
}
|
||||
|
||||
extra_args = {"StorageClass": settings.s3_storage_class}
|
||||
|
||||
# Add ACL if configured
|
||||
if settings.s3_acl:
|
||||
extra_args['ACL'] = settings.s3_acl
|
||||
|
||||
extra_args["ACL"] = settings.s3_acl
|
||||
|
||||
# Upload file
|
||||
s3_client.upload_file(
|
||||
file_path,
|
||||
settings.s3_bucket_name,
|
||||
file_path,
|
||||
settings.s3_bucket_name,
|
||||
s3_key,
|
||||
ExtraArgs=extra_args
|
||||
)
|
||||
|
||||
|
||||
# Generate URL to the file (useful for public files)
|
||||
# For private files, this is just a reference and won't be accessible directly
|
||||
s3_url = f"https://{settings.s3_bucket_name}.s3.{settings.aws_region}.amazonaws.com/{s3_key}"
|
||||
|
||||
logger.info(f"[{task_id}] Successfully uploaded {filename} to S3 bucket {settings.s3_bucket_name} at path {s3_key}")
|
||||
|
||||
logger.info(
|
||||
f"[{task_id}] Successfully uploaded {filename} to S3 bucket {settings.s3_bucket_name} at path {s3_key}"
|
||||
)
|
||||
log_task_progress(task_id, "upload_to_s3", "success", f"Uploaded to S3: {filename}", file_id=file_id)
|
||||
return {
|
||||
"status": "Completed",
|
||||
"file": file_path,
|
||||
"s3_bucket": settings.s3_bucket_name,
|
||||
"s3_key": s3_key,
|
||||
"s3_url": s3_url
|
||||
"s3_url": s3_url,
|
||||
}
|
||||
|
||||
|
||||
except ClientError as e:
|
||||
error_msg = f"Failed to upload {filename} to S3: {str(e)}"
|
||||
logger.error(f"[{task_id}] {error_msg}")
|
||||
log_task_progress(task_id, "upload_to_s3", "failure", error_msg, file_id=file_id)
|
||||
raise Exception(error_msg)
|
||||
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Error uploading {filename} to S3: {str(e)}"
|
||||
logger.error(f"[{task_id}] {error_msg}")
|
||||
|
||||
+62
-52
@@ -1,6 +1,7 @@
|
||||
"""
|
||||
Tests for upload tasks including OneDrive and S3.
|
||||
"""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
@@ -11,15 +12,16 @@ from app.tasks.upload_to_s3 import upload_to_s3
|
||||
@pytest.fixture
|
||||
def mock_settings():
|
||||
"""Mock settings for upload tests."""
|
||||
with patch('app.tasks.upload_to_onedrive.settings') as onedrive_settings, \
|
||||
patch('app.tasks.upload_to_s3.settings') as s3_settings:
|
||||
with patch("app.tasks.upload_to_onedrive.settings") as onedrive_settings, patch(
|
||||
"app.tasks.upload_to_s3.settings"
|
||||
) as s3_settings:
|
||||
# OneDrive settings
|
||||
onedrive_settings.onedrive_client_id = "test_client_id"
|
||||
onedrive_settings.onedrive_client_secret = "test_secret"
|
||||
onedrive_settings.onedrive_refresh_token = "test_refresh_token"
|
||||
onedrive_settings.onedrive_folder_path = "test_folder"
|
||||
onedrive_settings.onedrive_tenant_id = "common"
|
||||
|
||||
|
||||
# S3 settings
|
||||
s3_settings.s3_bucket_name = "test-bucket"
|
||||
s3_settings.aws_access_key_id = "test_access_key"
|
||||
@@ -28,27 +30,28 @@ def mock_settings():
|
||||
s3_settings.s3_folder_prefix = "documents"
|
||||
s3_settings.s3_storage_class = "STANDARD"
|
||||
s3_settings.s3_acl = None
|
||||
|
||||
|
||||
yield onedrive_settings, s3_settings
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_upload_to_onedrive_accepts_file_id(sample_text_file, mock_settings):
|
||||
"""Test that upload_to_onedrive accepts file_id parameter."""
|
||||
with patch('app.tasks.upload_to_onedrive.get_onedrive_token') as mock_token, \
|
||||
patch('app.tasks.upload_to_onedrive.create_upload_session') as mock_session, \
|
||||
patch('app.tasks.upload_to_onedrive.upload_large_file') as mock_upload, \
|
||||
patch('app.tasks.upload_to_onedrive.log_task_progress'):
|
||||
|
||||
with patch("app.tasks.upload_to_onedrive.get_onedrive_token") as mock_token, patch(
|
||||
"app.tasks.upload_to_onedrive.create_upload_session"
|
||||
) as mock_session, patch("app.tasks.upload_to_onedrive.upload_large_file") as mock_upload, patch(
|
||||
"app.tasks.upload_to_onedrive.log_task_progress"
|
||||
):
|
||||
|
||||
# Setup mocks
|
||||
mock_token.return_value = "test_access_token"
|
||||
mock_session.return_value = "https://upload.url"
|
||||
mock_upload.return_value = {"webUrl": "https://onedrive.test/file"}
|
||||
|
||||
|
||||
# Call with file_id parameter using apply() to simulate task execution
|
||||
# This bypasses Celery and calls the function directly
|
||||
result = upload_to_onedrive.apply(args=[sample_text_file], kwargs={'file_id': 42}).get()
|
||||
|
||||
result = upload_to_onedrive.apply(args=[sample_text_file], kwargs={"file_id": 42}).get()
|
||||
|
||||
assert result["status"] == "Completed"
|
||||
assert result["file_path"] == sample_text_file
|
||||
assert "onedrive_path" in result
|
||||
@@ -57,19 +60,20 @@ def test_upload_to_onedrive_accepts_file_id(sample_text_file, mock_settings):
|
||||
@pytest.mark.unit
|
||||
def test_upload_to_onedrive_without_file_id(sample_text_file, mock_settings):
|
||||
"""Test that upload_to_onedrive works without file_id parameter."""
|
||||
with patch('app.tasks.upload_to_onedrive.get_onedrive_token') as mock_token, \
|
||||
patch('app.tasks.upload_to_onedrive.create_upload_session') as mock_session, \
|
||||
patch('app.tasks.upload_to_onedrive.upload_large_file') as mock_upload, \
|
||||
patch('app.tasks.upload_to_onedrive.log_task_progress'):
|
||||
|
||||
with patch("app.tasks.upload_to_onedrive.get_onedrive_token") as mock_token, patch(
|
||||
"app.tasks.upload_to_onedrive.create_upload_session"
|
||||
) as mock_session, patch("app.tasks.upload_to_onedrive.upload_large_file") as mock_upload, patch(
|
||||
"app.tasks.upload_to_onedrive.log_task_progress"
|
||||
):
|
||||
|
||||
# Setup mocks
|
||||
mock_token.return_value = "test_access_token"
|
||||
mock_session.return_value = "https://upload.url"
|
||||
mock_upload.return_value = {"webUrl": "https://onedrive.test/file"}
|
||||
|
||||
|
||||
# Call without file_id parameter - should use default None
|
||||
result = upload_to_onedrive.apply(args=[sample_text_file]).get()
|
||||
|
||||
|
||||
assert result["status"] == "Completed"
|
||||
assert result["file_path"] == sample_text_file
|
||||
|
||||
@@ -77,17 +81,18 @@ def test_upload_to_onedrive_without_file_id(sample_text_file, mock_settings):
|
||||
@pytest.mark.unit
|
||||
def test_upload_to_s3_accepts_file_id(sample_text_file, mock_settings):
|
||||
"""Test that upload_to_s3 accepts file_id parameter."""
|
||||
with patch('app.tasks.upload_to_s3.boto3.client') as mock_boto_client, \
|
||||
patch('app.tasks.upload_to_s3.log_task_progress'):
|
||||
|
||||
with patch("app.tasks.upload_to_s3.boto3.client") as mock_boto_client, patch(
|
||||
"app.tasks.upload_to_s3.log_task_progress"
|
||||
):
|
||||
|
||||
# Setup mock S3 client
|
||||
mock_s3 = Mock()
|
||||
mock_boto_client.return_value = mock_s3
|
||||
mock_s3.upload_file.return_value = None
|
||||
|
||||
|
||||
# Call with file_id parameter using apply()
|
||||
result = upload_to_s3.apply(args=[sample_text_file], kwargs={'file_id': 99}).get()
|
||||
|
||||
result = upload_to_s3.apply(args=[sample_text_file], kwargs={"file_id": 99}).get()
|
||||
|
||||
assert result["status"] == "Completed"
|
||||
assert result["file"] == sample_text_file
|
||||
assert result["s3_bucket"] == "test-bucket"
|
||||
@@ -97,17 +102,18 @@ def test_upload_to_s3_accepts_file_id(sample_text_file, mock_settings):
|
||||
@pytest.mark.unit
|
||||
def test_upload_to_s3_without_file_id(sample_text_file, mock_settings):
|
||||
"""Test that upload_to_s3 works without file_id parameter."""
|
||||
with patch('app.tasks.upload_to_s3.boto3.client') as mock_boto_client, \
|
||||
patch('app.tasks.upload_to_s3.log_task_progress'):
|
||||
|
||||
with patch("app.tasks.upload_to_s3.boto3.client") as mock_boto_client, patch(
|
||||
"app.tasks.upload_to_s3.log_task_progress"
|
||||
):
|
||||
|
||||
# Setup mock S3 client
|
||||
mock_s3 = Mock()
|
||||
mock_boto_client.return_value = mock_s3
|
||||
mock_s3.upload_file.return_value = None
|
||||
|
||||
|
||||
# Call without file_id parameter using apply()
|
||||
result = upload_to_s3.apply(args=[sample_text_file]).get()
|
||||
|
||||
|
||||
assert result["status"] == "Completed"
|
||||
assert result["file"] == sample_text_file
|
||||
|
||||
@@ -115,60 +121,64 @@ def test_upload_to_s3_without_file_id(sample_text_file, mock_settings):
|
||||
@pytest.mark.unit
|
||||
def test_upload_to_onedrive_file_not_found(mock_settings):
|
||||
"""Test that upload_to_onedrive raises error for missing file."""
|
||||
with patch('app.tasks.upload_to_onedrive.log_task_progress'):
|
||||
with patch("app.tasks.upload_to_onedrive.log_task_progress"):
|
||||
with pytest.raises(FileNotFoundError):
|
||||
upload_to_onedrive.apply(args=["/nonexistent/file.pdf"], kwargs={'file_id': 1}).get()
|
||||
upload_to_onedrive.apply(args=["/nonexistent/file.pdf"], kwargs={"file_id": 1}).get()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_upload_to_s3_file_not_found(mock_settings):
|
||||
"""Test that upload_to_s3 raises error for missing file."""
|
||||
with patch('app.tasks.upload_to_s3.log_task_progress'):
|
||||
with patch("app.tasks.upload_to_s3.log_task_progress"):
|
||||
with pytest.raises(FileNotFoundError):
|
||||
upload_to_s3.apply(args=["/nonexistent/file.pdf"], kwargs={'file_id': 1}).get()
|
||||
upload_to_s3.apply(args=["/nonexistent/file.pdf"], kwargs={"file_id": 1}).get()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_upload_to_onedrive_logs_with_file_id(sample_text_file, mock_settings):
|
||||
"""Test that upload_to_onedrive properly logs with file_id."""
|
||||
with patch('app.tasks.upload_to_onedrive.get_onedrive_token') as mock_token, \
|
||||
patch('app.tasks.upload_to_onedrive.create_upload_session') as mock_session, \
|
||||
patch('app.tasks.upload_to_onedrive.upload_large_file') as mock_upload, \
|
||||
patch('app.tasks.upload_to_onedrive.log_task_progress') as mock_log:
|
||||
|
||||
with patch("app.tasks.upload_to_onedrive.get_onedrive_token") as mock_token, patch(
|
||||
"app.tasks.upload_to_onedrive.create_upload_session"
|
||||
) as mock_session, patch("app.tasks.upload_to_onedrive.upload_large_file") as mock_upload, patch(
|
||||
"app.tasks.upload_to_onedrive.log_task_progress"
|
||||
) as mock_log:
|
||||
|
||||
# Setup mocks
|
||||
mock_token.return_value = "test_access_token"
|
||||
mock_session.return_value = "https://upload.url"
|
||||
mock_upload.return_value = {"webUrl": "https://onedrive.test/file"}
|
||||
|
||||
|
||||
# Call with file_id
|
||||
upload_to_onedrive.apply(args=[sample_text_file], kwargs={'file_id': 123}).get()
|
||||
|
||||
upload_to_onedrive.apply(args=[sample_text_file], kwargs={"file_id": 123}).get()
|
||||
|
||||
# Verify log_task_progress was called with file_id
|
||||
assert mock_log.called
|
||||
# Check that at least one call included the file_id parameter
|
||||
calls_with_file_id = [call for call in mock_log.call_args_list
|
||||
if 'file_id' in call[1] and call[1]['file_id'] == 123]
|
||||
calls_with_file_id = [
|
||||
call for call in mock_log.call_args_list if "file_id" in call[1] and call[1]["file_id"] == 123
|
||||
]
|
||||
assert len(calls_with_file_id) > 0
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_upload_to_s3_logs_with_file_id(sample_text_file, mock_settings):
|
||||
"""Test that upload_to_s3 properly logs with file_id."""
|
||||
with patch('app.tasks.upload_to_s3.boto3.client') as mock_boto_client, \
|
||||
patch('app.tasks.upload_to_s3.log_task_progress') as mock_log:
|
||||
|
||||
with patch("app.tasks.upload_to_s3.boto3.client") as mock_boto_client, patch(
|
||||
"app.tasks.upload_to_s3.log_task_progress"
|
||||
) as mock_log:
|
||||
|
||||
# Setup mock S3 client
|
||||
mock_s3 = Mock()
|
||||
mock_boto_client.return_value = mock_s3
|
||||
mock_s3.upload_file.return_value = None
|
||||
|
||||
|
||||
# Call with file_id
|
||||
upload_to_s3.apply(args=[sample_text_file], kwargs={'file_id': 456}).get()
|
||||
|
||||
upload_to_s3.apply(args=[sample_text_file], kwargs={"file_id": 456}).get()
|
||||
|
||||
# Verify log_task_progress was called with file_id
|
||||
assert mock_log.called
|
||||
# Check that at least one call included the file_id parameter
|
||||
calls_with_file_id = [call for call in mock_log.call_args_list
|
||||
if 'file_id' in call[1] and call[1]['file_id'] == 456]
|
||||
calls_with_file_id = [
|
||||
call for call in mock_log.call_args_list if "file_id" in call[1] and call[1]["file_id"] == 456
|
||||
]
|
||||
assert len(calls_with_file_id) > 0
|
||||
|
||||
Reference in New Issue
Block a user