test: use sqlite3 session in test_remove_app_and_related_data_task (#38778)

This commit is contained in:
Asuka Minato 2026-07-13 12:02:30 +09:00 committed by GitHub
parent 7311f1ba6d
commit 4c1c92eb2f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,14 +1,20 @@
import logging
from collections.abc import Generator
from datetime import UTC, datetime
from unittest.mock import MagicMock, call, patch
from uuid import uuid4
import pytest
from agenton.compositor import CompositorSessionSnapshot
from sqlalchemy import delete, select
from sqlalchemy.orm import Session
from core.db.session_factory import session_factory
from graphon.enums import WorkflowExecutionStatus
from libs.archive_storage import ArchiveStorageNotConfiguredError
from models import AgentRuntimeSession, AgentRuntimeSessionOwnerType, AgentRuntimeSessionStatus
from models import AgentRuntimeSession, AgentRuntimeSessionOwnerType, AgentRuntimeSessionStatus, AppStar
from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom
from models.workflow import WorkflowArchiveLog
from tasks.remove_app_and_related_data_task import (
_cleanup_active_agent_runtime_sessions_for_app,
_delete_app_stars,
@ -99,9 +105,12 @@ class TestDeleteDraftVariableOffloadData:
class TestDeleteWorkflowArchiveLogs:
@pytest.mark.parametrize("sqlite_session", [(WorkflowArchiveLog,)], indirect=True)
@patch("tasks.remove_app_and_related_data_task._delete_records")
@patch("tasks.remove_app_and_related_data_task.db")
def test_delete_app_workflow_archive_logs_calls_delete_records(self, mock_db, mock_delete_records):
def test_delete_app_workflow_archive_logs_calls_delete_records(
self, mock_db, mock_delete_records, sqlite_session: Session
):
tenant_id = "tenant-1"
app_id = "app-1"
@ -113,16 +122,42 @@ class TestDeleteWorkflowArchiveLogs:
assert params == {"tenant_id": tenant_id, "app_id": app_id}
assert name == "workflow archive log"
mock_session = MagicMock()
archive_log = WorkflowArchiveLog(
tenant_id=str(uuid4()),
app_id=str(uuid4()),
workflow_id=str(uuid4()),
workflow_run_id=str(uuid4()),
created_by_role=CreatorUserRole.ACCOUNT,
created_by=str(uuid4()),
log_id=None,
log_created_at=None,
log_created_from=None,
run_version="1",
run_status=WorkflowExecutionStatus.SUCCEEDED,
run_triggered_from=WorkflowRunTriggeredFrom.APP_RUN,
run_error=None,
run_elapsed_time=0,
run_total_tokens=0,
run_total_steps=1,
run_created_at=datetime.now(UTC),
run_finished_at=datetime.now(UTC),
run_exceptions_count=0,
trigger_metadata=None,
)
sqlite_session.add(archive_log)
sqlite_session.commit()
delete_func(mock_session, "log-1")
delete_func(sqlite_session, archive_log.id)
sqlite_session.commit()
sqlite_session.expunge_all()
mock_session.execute.assert_called_once()
assert sqlite_session.get(WorkflowArchiveLog, archive_log.id) is None
class TestDeleteAppStars:
@pytest.mark.parametrize("sqlite_session", [(AppStar,)], indirect=True)
@patch("tasks.remove_app_and_related_data_task._delete_records")
def test_delete_app_stars_calls_delete_records(self, mock_delete_records):
def test_delete_app_stars_calls_delete_records(self, mock_delete_records, sqlite_session: Session):
tenant_id = "tenant-1"
app_id = "app-1"
@ -134,11 +169,15 @@ class TestDeleteAppStars:
assert params == {"tenant_id": tenant_id, "app_id": app_id}
assert name == "app star"
mock_session = MagicMock()
app_star = AppStar(tenant_id=str(uuid4()), app_id=str(uuid4()), account_id=str(uuid4()))
sqlite_session.add(app_star)
sqlite_session.commit()
delete_func(mock_session, "star-1")
delete_func(sqlite_session, app_star.id)
sqlite_session.commit()
sqlite_session.expunge_all()
mock_session.execute.assert_called_once()
assert sqlite_session.get(AppStar, app_star.id) is None
class TestDeleteArchivedWorkflowRunFiles: