mirror of
https://github.com/langgenius/dify.git
synced 2026-03-10 11:10:19 +08:00
test: migrate remove_app_and_related_data_task SQL tests to testcontainers (#32547)
Co-authored-by: KinomotoMio <200703522+KinomotoMio@users.noreply.github.com>
This commit is contained in:
木之本澪
GitHub
parent
02fef84d7f
commit
64296da7e7
@ -0,0 +1,224 @@
|
||||
import uuid
|
||||
from unittest.mock import ANY, call, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from core.db.session_factory import session_factory
|
||||
from core.variables.segments import StringSegment
|
||||
from core.variables.types import SegmentType
|
||||
from libs.datetime_utils import naive_utc_now
|
||||
from models import Tenant
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import App, UploadFile
|
||||
from models.workflow import WorkflowDraftVariable, WorkflowDraftVariableFile
|
||||
from tasks.remove_app_and_related_data_task import (
|
||||
_delete_draft_variable_offload_data,
|
||||
delete_draft_variables_batch,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cleanup_database(db_session_with_containers):
|
||||
db_session_with_containers.query(WorkflowDraftVariable).delete()
|
||||
db_session_with_containers.query(WorkflowDraftVariableFile).delete()
|
||||
db_session_with_containers.query(UploadFile).delete()
|
||||
db_session_with_containers.query(App).delete()
|
||||
db_session_with_containers.query(Tenant).delete()
|
||||
db_session_with_containers.commit()
|
||||
|
||||
|
||||
def _create_tenant_and_app(db_session_with_containers):
|
||||
tenant = Tenant(name=f"test_tenant_{uuid.uuid4()}")
|
||||
db_session_with_containers.add(tenant)
|
||||
db_session_with_containers.flush()
|
||||
|
||||
app = App(
|
||||
tenant_id=tenant.id,
|
||||
name=f"Test App for tenant {tenant.id}",
|
||||
mode="workflow",
|
||||
enable_site=True,
|
||||
enable_api=True,
|
||||
)
|
||||
db_session_with_containers.add(app)
|
||||
db_session_with_containers.commit()
|
||||
|
||||
return tenant, app
|
||||
|
||||
|
||||
def _create_draft_variables(
|
||||
db_session_with_containers,
|
||||
*,
|
||||
app_id: str,
|
||||
count: int,
|
||||
file_id_by_index: dict[int, str] | None = None,
|
||||
) -> list[WorkflowDraftVariable]:
|
||||
variables: list[WorkflowDraftVariable] = []
|
||||
file_id_by_index = file_id_by_index or {}
|
||||
|
||||
for i in range(count):
|
||||
variable = WorkflowDraftVariable.new_node_variable(
|
||||
app_id=app_id,
|
||||
node_id=f"node_{i}",
|
||||
name=f"var_{i}",
|
||||
value=StringSegment(value="test_value"),
|
||||
node_execution_id=str(uuid.uuid4()),
|
||||
file_id=file_id_by_index.get(i),
|
||||
)
|
||||
db_session_with_containers.add(variable)
|
||||
variables.append(variable)
|
||||
|
||||
db_session_with_containers.commit()
|
||||
return variables
|
||||
|
||||
|
||||
def _create_offload_data(db_session_with_containers, *, tenant_id: str, app_id: str, count: int):
|
||||
upload_files: list[UploadFile] = []
|
||||
variable_files: list[WorkflowDraftVariableFile] = []
|
||||
|
||||
for i in range(count):
|
||||
upload_file = UploadFile(
|
||||
tenant_id=tenant_id,
|
||||
storage_type="local",
|
||||
key=f"test/file-{uuid.uuid4()}-{i}.json",
|
||||
name=f"file-{i}.json",
|
||||
size=1024 + i,
|
||||
extension="json",
|
||||
mime_type="application/json",
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by=str(uuid.uuid4()),
|
||||
created_at=naive_utc_now(),
|
||||
used=False,
|
||||
)
|
||||
db_session_with_containers.add(upload_file)
|
||||
db_session_with_containers.flush()
|
||||
upload_files.append(upload_file)
|
||||
|
||||
variable_file = WorkflowDraftVariableFile(
|
||||
tenant_id=tenant_id,
|
||||
app_id=app_id,
|
||||
user_id=str(uuid.uuid4()),
|
||||
upload_file_id=upload_file.id,
|
||||
size=1024 + i,
|
||||
length=10 + i,
|
||||
value_type=SegmentType.STRING,
|
||||
)
|
||||
db_session_with_containers.add(variable_file)
|
||||
db_session_with_containers.flush()
|
||||
variable_files.append(variable_file)
|
||||
|
||||
db_session_with_containers.commit()
|
||||
|
||||
return {
|
||||
"upload_files": upload_files,
|
||||
"variable_files": variable_files,
|
||||
}
|
||||
|
||||
|
||||
class TestDeleteDraftVariablesBatch:
|
||||
def test_delete_draft_variables_batch_success(self, db_session_with_containers):
|
||||
"""Test successful deletion of draft variables in batches."""
|
||||
_, app1 = _create_tenant_and_app(db_session_with_containers)
|
||||
_, app2 = _create_tenant_and_app(db_session_with_containers)
|
||||
|
||||
_create_draft_variables(db_session_with_containers, app_id=app1.id, count=150)
|
||||
_create_draft_variables(db_session_with_containers, app_id=app2.id, count=100)
|
||||
|
||||
result = delete_draft_variables_batch(app1.id, batch_size=100)
|
||||
|
||||
assert result == 150
|
||||
app1_remaining = db_session_with_containers.query(WorkflowDraftVariable).where(
|
||||
WorkflowDraftVariable.app_id == app1.id
|
||||
)
|
||||
app2_remaining = db_session_with_containers.query(WorkflowDraftVariable).where(
|
||||
WorkflowDraftVariable.app_id == app2.id
|
||||
)
|
||||
assert app1_remaining.count() == 0
|
||||
assert app2_remaining.count() == 100
|
||||
|
||||
def test_delete_draft_variables_batch_empty_result(self, db_session_with_containers):
|
||||
"""Test deletion when no draft variables exist for the app."""
|
||||
result = delete_draft_variables_batch(str(uuid.uuid4()), 1000)
|
||||
|
||||
assert result == 0
|
||||
assert db_session_with_containers.query(WorkflowDraftVariable).count() == 0
|
||||
|
||||
@patch("tasks.remove_app_and_related_data_task._delete_draft_variable_offload_data")
|
||||
@patch("tasks.remove_app_and_related_data_task.logger")
|
||||
def test_delete_draft_variables_batch_logs_progress(
|
||||
self, mock_logger, mock_offload_cleanup, db_session_with_containers
|
||||
):
|
||||
"""Test that batch deletion logs progress correctly."""
|
||||
tenant, app = _create_tenant_and_app(db_session_with_containers)
|
||||
offload_data = _create_offload_data(db_session_with_containers, tenant_id=tenant.id, app_id=app.id, count=10)
|
||||
|
||||
file_ids = [variable_file.id for variable_file in offload_data["variable_files"]]
|
||||
file_id_by_index: dict[int, str] = {}
|
||||
for i in range(30):
|
||||
if i % 3 == 0:
|
||||
file_id_by_index[i] = file_ids[i // 3]
|
||||
_create_draft_variables(db_session_with_containers, app_id=app.id, count=30, file_id_by_index=file_id_by_index)
|
||||
|
||||
mock_offload_cleanup.return_value = len(file_id_by_index)
|
||||
|
||||
result = delete_draft_variables_batch(app.id, 50)
|
||||
|
||||
assert result == 30
|
||||
mock_offload_cleanup.assert_called_once()
|
||||
_, called_file_ids = mock_offload_cleanup.call_args.args
|
||||
assert {str(file_id) for file_id in called_file_ids} == {str(file_id) for file_id in file_id_by_index.values()}
|
||||
assert mock_logger.info.call_count == 2
|
||||
mock_logger.info.assert_any_call(ANY)
|
||||
|
||||
|
||||
class TestDeleteDraftVariableOffloadData:
|
||||
"""Test the Offload data cleanup functionality."""
|
||||
|
||||
@patch("extensions.ext_storage.storage")
|
||||
def test_delete_draft_variable_offload_data_success(self, mock_storage, db_session_with_containers):
|
||||
"""Test successful deletion of offload data."""
|
||||
tenant, app = _create_tenant_and_app(db_session_with_containers)
|
||||
offload_data = _create_offload_data(db_session_with_containers, tenant_id=tenant.id, app_id=app.id, count=3)
|
||||
file_ids = [variable_file.id for variable_file in offload_data["variable_files"]]
|
||||
upload_file_keys = [upload_file.key for upload_file in offload_data["upload_files"]]
|
||||
upload_file_ids = [upload_file.id for upload_file in offload_data["upload_files"]]
|
||||
|
||||
with session_factory.create_session() as session, session.begin():
|
||||
result = _delete_draft_variable_offload_data(session, file_ids)
|
||||
|
||||
assert result == 3
|
||||
expected_storage_calls = [call(storage_key) for storage_key in upload_file_keys]
|
||||
mock_storage.delete.assert_has_calls(expected_storage_calls, any_order=True)
|
||||
|
||||
remaining_var_files = db_session_with_containers.query(WorkflowDraftVariableFile).where(
|
||||
WorkflowDraftVariableFile.id.in_(file_ids)
|
||||
)
|
||||
remaining_upload_files = db_session_with_containers.query(UploadFile).where(UploadFile.id.in_(upload_file_ids))
|
||||
assert remaining_var_files.count() == 0
|
||||
assert remaining_upload_files.count() == 0
|
||||
|
||||
@patch("extensions.ext_storage.storage")
|
||||
@patch("tasks.remove_app_and_related_data_task.logging")
|
||||
def test_delete_draft_variable_offload_data_storage_failure(
|
||||
self, mock_logging, mock_storage, db_session_with_containers
|
||||
):
|
||||
"""Test handling of storage deletion failures."""
|
||||
tenant, app = _create_tenant_and_app(db_session_with_containers)
|
||||
offload_data = _create_offload_data(db_session_with_containers, tenant_id=tenant.id, app_id=app.id, count=2)
|
||||
file_ids = [variable_file.id for variable_file in offload_data["variable_files"]]
|
||||
storage_keys = [upload_file.key for upload_file in offload_data["upload_files"]]
|
||||
upload_file_ids = [upload_file.id for upload_file in offload_data["upload_files"]]
|
||||
|
||||
mock_storage.delete.side_effect = [Exception("Storage error"), None]
|
||||
|
||||
with session_factory.create_session() as session, session.begin():
|
||||
result = _delete_draft_variable_offload_data(session, file_ids)
|
||||
|
||||
assert result == 1
|
||||
mock_logging.exception.assert_called_once_with("Failed to delete storage object %s", storage_keys[0])
|
||||
|
||||
remaining_var_files = db_session_with_containers.query(WorkflowDraftVariableFile).where(
|
||||
WorkflowDraftVariableFile.id.in_(file_ids)
|
||||
)
|
||||
remaining_upload_files = db_session_with_containers.query(UploadFile).where(UploadFile.id.in_(upload_file_ids))
|
||||
assert remaining_var_files.count() == 0
|
||||
assert remaining_upload_files.count() == 0
|
||||
@ -1,4 +1,4 @@
|
||||
from unittest.mock import ANY, MagicMock, call, patch
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
|
||||
import pytest
|
||||
|
||||
@ -14,124 +14,6 @@ from tasks.remove_app_and_related_data_task import (
|
||||
|
||||
|
||||
class TestDeleteDraftVariablesBatch:
|
||||
@patch("tasks.remove_app_and_related_data_task._delete_draft_variable_offload_data")
|
||||
@patch("tasks.remove_app_and_related_data_task.session_factory")
|
||||
def test_delete_draft_variables_batch_success(self, mock_sf, mock_offload_cleanup):
|
||||
"""Test successful deletion of draft variables in batches."""
|
||||
app_id = "test-app-id"
|
||||
batch_size = 100
|
||||
|
||||
# Mock session via session_factory
|
||||
mock_session = MagicMock()
|
||||
mock_context_manager = MagicMock()
|
||||
mock_context_manager.__enter__.return_value = mock_session
|
||||
mock_context_manager.__exit__.return_value = None
|
||||
mock_sf.create_session.return_value = mock_context_manager
|
||||
|
||||
# Mock two batches of results, then empty
|
||||
batch1_data = [(f"var-{i}", f"file-{i}" if i % 2 == 0 else None) for i in range(100)]
|
||||
batch2_data = [(f"var-{i}", f"file-{i}" if i % 3 == 0 else None) for i in range(100, 150)]
|
||||
|
||||
batch1_ids = [row[0] for row in batch1_data]
|
||||
batch1_file_ids = [row[1] for row in batch1_data if row[1] is not None]
|
||||
|
||||
batch2_ids = [row[0] for row in batch2_data]
|
||||
batch2_file_ids = [row[1] for row in batch2_data if row[1] is not None]
|
||||
|
||||
# Setup side effects for execute calls in the correct order:
|
||||
# 1. SELECT (returns batch1_data with id, file_id)
|
||||
# 2. DELETE (returns result with rowcount=100)
|
||||
# 3. SELECT (returns batch2_data)
|
||||
# 4. DELETE (returns result with rowcount=50)
|
||||
# 5. SELECT (returns empty, ends loop)
|
||||
|
||||
# Create mock results with actual integer rowcount attributes
|
||||
class MockResult:
|
||||
def __init__(self, rowcount):
|
||||
self.rowcount = rowcount
|
||||
|
||||
# First SELECT result
|
||||
select_result1 = MagicMock()
|
||||
select_result1.__iter__.return_value = iter(batch1_data)
|
||||
|
||||
# First DELETE result
|
||||
delete_result1 = MockResult(rowcount=100)
|
||||
|
||||
# Second SELECT result
|
||||
select_result2 = MagicMock()
|
||||
select_result2.__iter__.return_value = iter(batch2_data)
|
||||
|
||||
# Second DELETE result
|
||||
delete_result2 = MockResult(rowcount=50)
|
||||
|
||||
# Third SELECT result (empty, ends loop)
|
||||
select_result3 = MagicMock()
|
||||
select_result3.__iter__.return_value = iter([])
|
||||
|
||||
# Configure side effects in the correct order
|
||||
mock_session.execute.side_effect = [
|
||||
select_result1, # First SELECT
|
||||
delete_result1, # First DELETE
|
||||
select_result2, # Second SELECT
|
||||
delete_result2, # Second DELETE
|
||||
select_result3, # Third SELECT (empty)
|
||||
]
|
||||
|
||||
# Mock offload data cleanup
|
||||
mock_offload_cleanup.side_effect = [len(batch1_file_ids), len(batch2_file_ids)]
|
||||
|
||||
# Execute the function
|
||||
result = delete_draft_variables_batch(app_id, batch_size)
|
||||
|
||||
# Verify the result
|
||||
assert result == 150
|
||||
|
||||
# Verify database calls
|
||||
assert mock_session.execute.call_count == 5 # 3 selects + 2 deletes
|
||||
|
||||
# Verify offload cleanup was called for both batches with file_ids
|
||||
expected_offload_calls = [call(mock_session, batch1_file_ids), call(mock_session, batch2_file_ids)]
|
||||
mock_offload_cleanup.assert_has_calls(expected_offload_calls)
|
||||
|
||||
# Simplified verification - check that the right number of calls were made
|
||||
# and that the SQL queries contain the expected patterns
|
||||
actual_calls = mock_session.execute.call_args_list
|
||||
for i, actual_call in enumerate(actual_calls):
|
||||
sql_text = str(actual_call[0][0])
|
||||
normalized = " ".join(sql_text.split())
|
||||
if i % 2 == 0: # SELECT calls (even indices: 0, 2, 4)
|
||||
assert "SELECT id, file_id FROM workflow_draft_variables" in normalized
|
||||
assert "WHERE app_id = :app_id" in normalized
|
||||
assert "LIMIT :batch_size" in normalized
|
||||
else: # DELETE calls (odd indices: 1, 3)
|
||||
assert "DELETE FROM workflow_draft_variables" in normalized
|
||||
assert "WHERE id IN :ids" in normalized
|
||||
|
||||
@patch("tasks.remove_app_and_related_data_task._delete_draft_variable_offload_data")
|
||||
@patch("tasks.remove_app_and_related_data_task.session_factory")
|
||||
def test_delete_draft_variables_batch_empty_result(self, mock_sf, mock_offload_cleanup):
|
||||
"""Test deletion when no draft variables exist for the app."""
|
||||
app_id = "nonexistent-app-id"
|
||||
batch_size = 1000
|
||||
|
||||
# Mock session via session_factory
|
||||
mock_session = MagicMock()
|
||||
mock_context_manager = MagicMock()
|
||||
mock_context_manager.__enter__.return_value = mock_session
|
||||
mock_context_manager.__exit__.return_value = None
|
||||
mock_sf.create_session.return_value = mock_context_manager
|
||||
|
||||
# Mock empty result
|
||||
empty_result = MagicMock()
|
||||
empty_result.__iter__.return_value = iter([])
|
||||
mock_session.execute.return_value = empty_result
|
||||
|
||||
result = delete_draft_variables_batch(app_id, batch_size)
|
||||
|
||||
assert result == 0
|
||||
assert mock_session.execute.call_count == 1 # Only one select query
|
||||
mock_offload_cleanup.assert_not_called() # No files to clean up
|
||||
|
||||
def test_delete_draft_variables_batch_invalid_batch_size(self):
|
||||
"""Test that invalid batch size raises ValueError."""
|
||||
app_id = "test-app-id"
|
||||
@ -142,66 +24,6 @@ class TestDeleteDraftVariablesBatch:
|
||||
with pytest.raises(ValueError, match="batch_size must be positive"):
|
||||
delete_draft_variables_batch(app_id, 0)
|
||||
|
||||
@patch("tasks.remove_app_and_related_data_task._delete_draft_variable_offload_data")
|
||||
@patch("tasks.remove_app_and_related_data_task.session_factory")
|
||||
@patch("tasks.remove_app_and_related_data_task.logger")
|
||||
def test_delete_draft_variables_batch_logs_progress(self, mock_logging, mock_sf, mock_offload_cleanup):
|
||||
"""Test that batch deletion logs progress correctly."""
|
||||
app_id = "test-app-id"
|
||||
batch_size = 50
|
||||
|
||||
# Mock session via session_factory
|
||||
mock_session = MagicMock()
|
||||
mock_context_manager = MagicMock()
|
||||
mock_context_manager.__enter__.return_value = mock_session
|
||||
mock_context_manager.__exit__.return_value = None
|
||||
mock_sf.create_session.return_value = mock_context_manager
|
||||
|
||||
# Mock one batch then empty
|
||||
batch_data = [(f"var-{i}", f"file-{i}" if i % 3 == 0 else None) for i in range(30)]
|
||||
batch_ids = [row[0] for row in batch_data]
|
||||
batch_file_ids = [row[1] for row in batch_data if row[1] is not None]
|
||||
|
||||
# Create properly configured mocks
|
||||
select_result = MagicMock()
|
||||
select_result.__iter__.return_value = iter(batch_data)
|
||||
|
||||
# Create simple object with rowcount attribute
|
||||
class MockResult:
|
||||
def __init__(self, rowcount):
|
||||
self.rowcount = rowcount
|
||||
|
||||
delete_result = MockResult(rowcount=30)
|
||||
|
||||
empty_result = MagicMock()
|
||||
empty_result.__iter__.return_value = iter([])
|
||||
|
||||
mock_session.execute.side_effect = [
|
||||
# Select query result
|
||||
select_result,
|
||||
# Delete query result
|
||||
delete_result,
|
||||
# Empty select result (end condition)
|
||||
empty_result,
|
||||
]
|
||||
|
||||
# Mock offload cleanup
|
||||
mock_offload_cleanup.return_value = len(batch_file_ids)
|
||||
|
||||
result = delete_draft_variables_batch(app_id, batch_size)
|
||||
|
||||
assert result == 30
|
||||
|
||||
# Verify offload cleanup was called with file_ids
|
||||
if batch_file_ids:
|
||||
mock_offload_cleanup.assert_called_once_with(mock_session, batch_file_ids)
|
||||
|
||||
# Verify logging calls
|
||||
assert mock_logging.info.call_count == 2
|
||||
mock_logging.info.assert_any_call(
|
||||
ANY # click.style call
|
||||
)
|
||||
|
||||
@patch("tasks.remove_app_and_related_data_task.delete_draft_variables_batch")
|
||||
def test_delete_draft_variables_calls_batch_function(self, mock_batch_delete):
|
||||
"""Test that _delete_draft_variables calls the batch function correctly."""
|
||||
@ -218,58 +40,6 @@ class TestDeleteDraftVariablesBatch:
|
||||
class TestDeleteDraftVariableOffloadData:
|
||||
"""Test the Offload data cleanup functionality."""
|
||||
|
||||
@patch("extensions.ext_storage.storage")
|
||||
def test_delete_draft_variable_offload_data_success(self, mock_storage):
|
||||
"""Test successful deletion of offload data."""
|
||||
|
||||
# Mock connection
|
||||
mock_conn = MagicMock()
|
||||
file_ids = ["file-1", "file-2", "file-3"]
|
||||
|
||||
# Mock query results: (variable_file_id, storage_key, upload_file_id)
|
||||
query_results = [
|
||||
("file-1", "storage/key/1", "upload-1"),
|
||||
("file-2", "storage/key/2", "upload-2"),
|
||||
("file-3", "storage/key/3", "upload-3"),
|
||||
]
|
||||
|
||||
mock_result = MagicMock()
|
||||
mock_result.__iter__.return_value = iter(query_results)
|
||||
mock_conn.execute.return_value = mock_result
|
||||
|
||||
# Execute function
|
||||
result = _delete_draft_variable_offload_data(mock_conn, file_ids)
|
||||
|
||||
# Verify return value
|
||||
assert result == 3
|
||||
|
||||
# Verify storage deletion calls
|
||||
expected_storage_calls = [call("storage/key/1"), call("storage/key/2"), call("storage/key/3")]
|
||||
mock_storage.delete.assert_has_calls(expected_storage_calls, any_order=True)
|
||||
|
||||
# Verify database calls - should be 3 calls total
|
||||
assert mock_conn.execute.call_count == 3
|
||||
|
||||
# Verify the queries were called
|
||||
actual_calls = mock_conn.execute.call_args_list
|
||||
|
||||
# First call should be the SELECT query
|
||||
select_call_sql = " ".join(str(actual_calls[0][0][0]).split())
|
||||
assert "SELECT wdvf.id, uf.key, uf.id as upload_file_id" in select_call_sql
|
||||
assert "FROM workflow_draft_variable_files wdvf" in select_call_sql
|
||||
assert "JOIN upload_files uf ON wdvf.upload_file_id = uf.id" in select_call_sql
|
||||
assert "WHERE wdvf.id IN :file_ids" in select_call_sql
|
||||
|
||||
# Second call should be DELETE upload_files
|
||||
delete_upload_call_sql = " ".join(str(actual_calls[1][0][0]).split())
|
||||
assert "DELETE FROM upload_files" in delete_upload_call_sql
|
||||
assert "WHERE id IN :upload_file_ids" in delete_upload_call_sql
|
||||
|
||||
# Third call should be DELETE workflow_draft_variable_files
|
||||
delete_variable_files_call_sql = " ".join(str(actual_calls[2][0][0]).split())
|
||||
assert "DELETE FROM workflow_draft_variable_files" in delete_variable_files_call_sql
|
||||
assert "WHERE id IN :file_ids" in delete_variable_files_call_sql
|
||||
|
||||
def test_delete_draft_variable_offload_data_empty_file_ids(self):
|
||||
"""Test handling of empty file_ids list."""
|
||||
mock_conn = MagicMock()
|
||||
@ -279,38 +49,6 @@ class TestDeleteDraftVariableOffloadData:
|
||||
assert result == 0
|
||||
mock_conn.execute.assert_not_called()
|
||||
|
||||
@patch("extensions.ext_storage.storage")
|
||||
@patch("tasks.remove_app_and_related_data_task.logging")
|
||||
def test_delete_draft_variable_offload_data_storage_failure(self, mock_logging, mock_storage):
|
||||
"""Test handling of storage deletion failures."""
|
||||
mock_conn = MagicMock()
|
||||
file_ids = ["file-1", "file-2"]
|
||||
|
||||
# Mock query results
|
||||
query_results = [
|
||||
("file-1", "storage/key/1", "upload-1"),
|
||||
("file-2", "storage/key/2", "upload-2"),
|
||||
]
|
||||
|
||||
mock_result = MagicMock()
|
||||
mock_result.__iter__.return_value = iter(query_results)
|
||||
mock_conn.execute.return_value = mock_result
|
||||
|
||||
# Make storage.delete fail for the first file
|
||||
mock_storage.delete.side_effect = [Exception("Storage error"), None]
|
||||
|
||||
# Execute function
|
||||
result = _delete_draft_variable_offload_data(mock_conn, file_ids)
|
||||
|
||||
# Should still return 2 (both files processed, even if one storage delete failed)
|
||||
assert result == 1 # Only one storage deletion succeeded
|
||||
|
||||
# Verify warning was logged
|
||||
mock_logging.exception.assert_called_once_with("Failed to delete storage object %s", "storage/key/1")
|
||||
|
||||
# Verify both database cleanup calls still happened
|
||||
assert mock_conn.execute.call_count == 3
|
||||
|
||||
@patch("tasks.remove_app_and_related_data_task.logging")
|
||||
def test_delete_draft_variable_offload_data_database_failure(self, mock_logging):
|
||||
"""Test handling of database operation failures."""
|
||||
|
||||
Loading…
Reference in New Issue
Block a user