From 4c1c92eb2f12c3aec668be55e7dda64ec0826b14 Mon Sep 17 00:00:00 2001
From: Asuka Minato <i@asukaminato.eu.org>
Date: Mon, 13 Jul 2026 12:02:30 +0900
Subject: [PATCH] test: use sqlite3 session in
 test_remove_app_and_related_data_task (#38778)

---
 .../test_remove_app_and_related_data_task.py  | 57 ++++++++++++++++---
 1 file changed, 48 insertions(+), 9 deletions(-)

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 1e4b37612ba..7d2705447db 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,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: