add clean sandbox workflow runs

api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py

@@ -4,6 +4,7 @@ from datetime import UTC, datetime
 from unittest.mock import Mock, patch
 
 import pytest
+from sqlalchemy.dialects import postgresql
 from sqlalchemy.orm import Session, sessionmaker
 
 from core.workflow.enums import WorkflowExecutionStatus
@@ -104,6 +105,42 @@ class TestDifyAPISQLAlchemyWorkflowRunRepository:
         return pause
 
 
+class TestGetRunsBatchByTimeRange(TestDifyAPISQLAlchemyWorkflowRunRepository):
+    def test_get_runs_batch_by_time_range_filters_terminal_statuses(
+        self, repository: DifyAPISQLAlchemyWorkflowRunRepository, mock_session: Mock
+    ):
+        scalar_result = Mock()
+        scalar_result.all.return_value = []
+        mock_session.scalars.return_value = scalar_result
+
+        repository.get_runs_batch_by_time_range(
+            start_after=None,
+            end_before=datetime(2024, 1, 1),
+            last_seen=None,
+            batch_size=50,
+        )
+
+        stmt = mock_session.scalars.call_args[0][0]
+        compiled_sql = str(
+            stmt.compile(
+                dialect=postgresql.dialect(),
+                compile_kwargs={"literal_binds": True},
+            )
+        )
+
+        assert "workflow_runs.status" in compiled_sql
+        for status in (
+            WorkflowExecutionStatus.SUCCEEDED,
+            WorkflowExecutionStatus.FAILED,
+            WorkflowExecutionStatus.STOPPED,
+            WorkflowExecutionStatus.PARTIAL_SUCCEEDED,
+        ):
+            assert f"'{status.value}'" in compiled_sql
+
+        assert "'running'" not in compiled_sql
+        assert "'paused'" not in compiled_sql
+
+
 class TestCreateWorkflowPause(TestDifyAPISQLAlchemyWorkflowRunRepository):
     """Test create_workflow_pause method."""
 
@@ -181,6 +218,31 @@ class TestCreateWorkflowPause(TestDifyAPISQLAlchemyWorkflowRunRepository):
             )
 
 
+class TestDeleteRunsWithRelated(TestDifyAPISQLAlchemyWorkflowRunRepository):
+    def test_uses_trigger_log_repository(self, repository: DifyAPISQLAlchemyWorkflowRunRepository, mock_session: Mock):
+        node_ids_result = Mock()
+        node_ids_result.all.return_value = []
+        pause_ids_result = Mock()
+        pause_ids_result.all.return_value = []
+        mock_session.scalars.side_effect = [node_ids_result, pause_ids_result]
+
+        # app_logs delete, runs delete
+        mock_session.execute.side_effect = [Mock(rowcount=0), Mock(rowcount=1)]
+
+        fake_trigger_repo = Mock()
+        fake_trigger_repo.delete_by_run_ids.return_value = 3
+
+        with patch(
+            "repositories.sqlalchemy_api_workflow_run_repository.SQLAlchemyWorkflowTriggerLogRepository",
+            return_value=fake_trigger_repo,
+        ):
+            counts = repository.delete_runs_with_related(["run-1"])
+
+        fake_trigger_repo.delete_by_run_ids.assert_called_once_with(["run-1"])
+        assert counts["trigger_logs"] == 3
+        assert counts["runs"] == 1
+
+
 class TestResumeWorkflowPause(TestDifyAPISQLAlchemyWorkflowRunRepository):
     """Test resume_workflow_pause method."""
 

api/tests/unit_tests/repositories/test_sqlalchemy_workflow_trigger_log_repository.py

@@ -0,0 +1,31 @@
+from unittest.mock import Mock
+
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.orm import Session
+
+from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository
+
+
+def test_delete_by_run_ids_executes_delete():
+    session = Mock(spec=Session)
+    session.execute.return_value = Mock(rowcount=2)
+    repo = SQLAlchemyWorkflowTriggerLogRepository(session)
+
+    deleted = repo.delete_by_run_ids(["run-1", "run-2"])
+
+    stmt = session.execute.call_args[0][0]
+    compiled_sql = str(stmt.compile(dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True}))
+    assert "workflow_trigger_logs" in compiled_sql
+    assert "'run-1'" in compiled_sql
+    assert "'run-2'" in compiled_sql
+    assert deleted == 2
+
+
+def test_delete_by_run_ids_empty_short_circuits():
+    session = Mock(spec=Session)
+    repo = SQLAlchemyWorkflowTriggerLogRepository(session)
+
+    deleted = repo.delete_by_run_ids([])
+
+    session.execute.assert_not_called()
+    assert deleted == 0

api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py

@@ -0,0 +1,175 @@
+import datetime
+from typing import Any
+
+import pytest
+
+from services import clear_free_plan_expired_workflow_run_logs as cleanup_module
+from services.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup
+
+
+class FakeRun:
+    def __init__(self, run_id: str, tenant_id: str, created_at: datetime.datetime) -> None:
+        self.id = run_id
+        self.tenant_id = tenant_id
+        self.created_at = created_at
+
+
+class FakeRepo:
+    def __init__(self, batches: list[list[FakeRun]], delete_result: dict[str, int] | None = None) -> None:
+        self.batches = batches
+        self.call_idx = 0
+        self.deleted: list[list[str]] = []
+        self.delete_result = delete_result or {
+            "runs": 0,
+            "node_executions": 0,
+            "offloads": 0,
+            "app_logs": 0,
+            "trigger_logs": 0,
+            "pauses": 0,
+            "pause_reasons": 0,
+        }
+
+    def get_runs_batch_by_time_range(
+        self,
+        start_after: datetime.datetime | None,
+        end_before: datetime.datetime,
+        last_seen: tuple[datetime.datetime, str] | None,
+        batch_size: int,
+    ) -> list[FakeRun]:
+        if self.call_idx >= len(self.batches):
+            return []
+        batch = self.batches[self.call_idx]
+        self.call_idx += 1
+        return batch
+
+    def delete_runs_with_related(self, run_ids: list[str]) -> dict[str, int]:
+        self.deleted.append(list(run_ids))
+        result = self.delete_result.copy()
+        result["runs"] = len(run_ids)
+        return result
+
+
+def test_filter_free_tenants_billing_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
+    cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([]))
+
+    monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", False)
+
+    def fail_bulk(_: list[str]) -> dict[str, dict[str, Any]]:
+        raise RuntimeError("should not call")
+
+    monkeypatch.setattr(cleanup_module.BillingService, "get_info_bulk", staticmethod(fail_bulk))
+
+    tenants = {"t1", "t2"}
+    free = cleanup._filter_free_tenants(tenants)
+
+    assert free == tenants
+
+
+def test_filter_free_tenants_bulk_mixed(monkeypatch: pytest.MonkeyPatch) -> None:
+    cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([]))
+
+    monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True)
+    cleanup.billing_cache["t_free"] = cleanup_module.CloudPlan.SANDBOX
+    cleanup.billing_cache["t_paid"] = cleanup_module.CloudPlan.TEAM
+    monkeypatch.setattr(
+        cleanup_module.BillingService,
+        "get_info_bulk",
+        staticmethod(lambda tenant_ids: dict.fromkeys(tenant_ids, "sandbox")),
+    )
+
+    free = cleanup._filter_free_tenants({"t_free", "t_paid", "t_missing"})
+
+    assert free == {"t_free", "t_missing"}
+
+
+def test_filter_free_tenants_bulk_failure(monkeypatch: pytest.MonkeyPatch) -> None:
+    cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([]))
+
+    monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True)
+    monkeypatch.setattr(
+        cleanup_module.BillingService,
+        "get_info_bulk",
+        staticmethod(lambda tenant_ids: (_ for _ in ()).throw(RuntimeError("boom"))),
+    )
+
+    free = cleanup._filter_free_tenants({"t1", "t2"})
+
+    assert free == set()
+
+
+def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None:
+    cutoff = datetime.datetime.now()
+    repo = FakeRepo(
+        batches=[
+            [
+                FakeRun("run-free", "t_free", cutoff),
+                FakeRun("run-paid", "t_paid", cutoff),
+            ]
+        ]
+    )
+    cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=repo)
+
+    monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True)
+    cleanup.billing_cache["t_free"] = cleanup_module.CloudPlan.SANDBOX
+    cleanup.billing_cache["t_paid"] = cleanup_module.CloudPlan.TEAM
+    monkeypatch.setattr(
+        cleanup_module.BillingService,
+        "get_info_bulk",
+        staticmethod(lambda tenant_ids: dict.fromkeys(tenant_ids, "sandbox")),
+    )
+
+    cleanup.run()
+
+    assert repo.deleted == [["run-free"]]
+
+
+def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch) -> None:
+    cutoff = datetime.datetime.now()
+    repo = FakeRepo(batches=[[FakeRun("run-paid", "t_paid", cutoff)]])
+    cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=repo)
+
+    monkeypatch.setattr(cleanup_module.dify_config, "BILLING_ENABLED", True)
+    monkeypatch.setattr(
+        cleanup_module.BillingService,
+        "get_info_bulk",
+        staticmethod(lambda tenant_ids: dict.fromkeys(tenant_ids, "team")),
+    )
+
+    cleanup.run()
+
+    assert repo.deleted == []
+
+
+def test_run_exits_on_empty_batch() -> None:
+    cleanup = WorkflowRunCleanup(days=30, batch_size=10, repo=FakeRepo([]))
+
+    cleanup.run()
+
+
+def test_between_sets_window_bounds() -> None:
+    start_after = datetime.datetime(2024, 5, 1, 0, 0, 0)
+    end_before = datetime.datetime(2024, 6, 1, 0, 0, 0)
+    cleanup = WorkflowRunCleanup(
+        days=30, batch_size=10, start_after=start_after, end_before=end_before, repo=FakeRepo([])
+    )
+
+    assert cleanup.window_start == start_after
+    assert cleanup.window_end == end_before
+
+
+def test_between_requires_both_boundaries() -> None:
+    with pytest.raises(ValueError):
+        WorkflowRunCleanup(
+            days=30, batch_size=10, start_after=datetime.datetime.now(), end_before=None, repo=FakeRepo([])
+        )
+    with pytest.raises(ValueError):
+        WorkflowRunCleanup(
+            days=30, batch_size=10, start_after=None, end_before=datetime.datetime.now(), repo=FakeRepo([])
+        )
+
+
+def test_between_requires_end_after_start() -> None:
+    start_after = datetime.datetime(2024, 6, 1, 0, 0, 0)
+    end_before = datetime.datetime(2024, 5, 1, 0, 0, 0)
+    with pytest.raises(ValueError):
+        WorkflowRunCleanup(days=30, batch_size=10, start_after=start_after, end_before=end_before, repo=FakeRepo([]))