From 64296da7e7eac635b6475f0b7085cd6b64f47308 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9C=A8=E4=B9=8B=E6=9C=AC=E6=BE=AA?= Date: Wed, 25 Feb 2026 04:12:23 +0800 Subject: [PATCH] test: migrate remove_app_and_related_data_task SQL tests to testcontainers (#32547) Co-authored-by: KinomotoMio <200703522+KinomotoMio@users.noreply.github.com> --- .../test_remove_app_and_related_data_task.py | 224 +++++++++++++++ .../test_remove_app_and_related_data_task.py | 264 +----------------- 2 files changed, 225 insertions(+), 263 deletions(-) create mode 100644 api/tests/test_containers_integration_tests/tasks/test_remove_app_and_related_data_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/test_containers_integration_tests/tasks/test_remove_app_and_related_data_task.py new file mode 100644 index 0000000000..7ac9573ab7 --- /dev/null +++ b/api/tests/test_containers_integration_tests/tasks/test_remove_app_and_related_data_task.py @@ -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 diff --git a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py index 2b11e42cd5..0ed4ca05fa 100644 --- a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py +++ b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py @@ -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."""