test: use SQLite sessions in unit misc (#39123)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/controllers/console/datasets/datasets.py

@@ -875,7 +875,7 @@ class DatasetIndexingEstimateApi(Resource):
                 file_details = session.scalars(
                     select(UploadFile).where(UploadFile.tenant_id == current_tenant_id, UploadFile.id.in_(file_ids))
                 ).all()
-                if file_details is None:
+                if not file_details:
                     raise NotFound("File not found.")
 
                 if file_details:

api/tests/unit_tests/controllers/console/test_wraps.py

@@ -1,3 +1,4 @@
+from types import SimpleNamespace
 from typing import override
 from unittest.mock import MagicMock, patch
 
@@ -5,6 +6,8 @@ import pytest
 from flask import Flask, request
 from flask_login import LoginManager, UserMixin
 from pydantic import BaseModel
+from sqlalchemy import select
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import HTTPException
 
 from controllers.common.wraps import _extract_resource_id
@@ -34,6 +37,7 @@ from controllers.console.wraps import (
 )
 from models import Account
 from models.account import AccountStatus, TenantAccountRole
+from models.dataset import RateLimitLog
 from services.feature_service import LicenseStatus
 
 
@@ -657,8 +661,7 @@ class TestRateLimiting:
     """Test rate limiting decorator"""
 
     @patch("controllers.console.wraps.redis_client")
-    @patch("controllers.console.wraps.db")
-    def test_should_allow_requests_within_rate_limit(self, mock_db: MagicMock, mock_redis: MagicMock):
+    def test_should_allow_requests_within_rate_limit(self, mock_redis: MagicMock):
         """Test that requests within rate limit are allowed"""
         # Arrange
         mock_rate_limit = MagicMock()
@@ -685,8 +688,13 @@ class TestRateLimiting:
         mock_redis.zremrangebyscore.assert_called_once()
 
     @patch("controllers.console.wraps.redis_client")
-    @patch("controllers.console.wraps.db")
-    def test_should_reject_requests_over_rate_limit(self, mock_db: MagicMock, mock_redis: MagicMock):
+    @pytest.mark.parametrize("sqlite_session", [(RateLimitLog,)], indirect=True)
+    def test_should_reject_requests_over_rate_limit(
+        self,
+        mock_redis: MagicMock,
+        sqlite_session: Session,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
         """Test that requests over rate limit are rejected and logged"""
         # Arrange
         app = create_app_with_login()
@@ -696,8 +704,7 @@ class TestRateLimiting:
         mock_rate_limit.subscription_plan = "pro"
         mock_redis.zcard.return_value = 11  # Over limit
 
-        mock_session = MagicMock()
-        mock_db.session = mock_session
+        monkeypatch.setattr("controllers.console.wraps.db", SimpleNamespace(session=sqlite_session))
 
         @cloud_edition_billing_rate_limit_check("knowledge")
         def knowledge_request():
@@ -719,9 +726,11 @@ class TestRateLimiting:
                     assert exc_info.value.code == 403
                     assert "rate limit" in str(exc_info.value.description)
 
-                    # Verify rate limit log was created
-                    mock_session.add.assert_called_once()
-                    mock_session.commit.assert_called_once()
+                    rate_limit_log = sqlite_session.scalar(select(RateLimitLog))
+                    assert rate_limit_log is not None
+                    assert rate_limit_log.tenant_id == "tenant123"
+                    assert rate_limit_log.subscription_plan == "pro"
+                    assert rate_limit_log.operation == "knowledge"
 
 
 class TestCloudUtmRecord:

api/tests/unit_tests/core/agent/test_cot_agent_runner.py

@@ -1,8 +1,11 @@
 import json
+from collections.abc import Iterator
 from unittest.mock import MagicMock
 
 import pytest
 from pytest_mock import MockerFixture
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
 
 from core.agent.cot_agent_runner import CotAgentRunner
 from core.agent.entities import AgentScratchpadUnit
@@ -26,7 +29,7 @@ class DummyRunner(CotAgentRunner):
 
 
 @pytest.fixture
-def runner(mocker: MockerFixture):
+def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[DummyRunner]:
     # Prevent BaseAgentRunner __init__ from hitting database
     mocker.patch(
         "core.agent.base_agent_runner.BaseAgentRunner.organize_agent_history",
@@ -81,9 +84,12 @@ def runner(mocker: MockerFixture):
     runner.agent_callback = None
     runner.memory = None
     runner.history_prompt_messages = []
-    runner.session = MagicMock()
+    runner.session = Session(sqlite_engine)
 
-    return runner
+    try:
+        yield runner
+    finally:
+        runner.session.close()
 
 
 class TestFillInputs:

api/tests/unit_tests/core/agent/test_fc_agent_runner.py

@@ -1,9 +1,12 @@
 import json
+from collections.abc import Iterator
 from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
 from pytest_mock import MockerFixture
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
 
 from core.agent.errors import AgentMaxIterationError
 from core.agent.fc_agent_runner import FunctionCallAgentRunner
@@ -69,7 +72,7 @@ class DummyResult:
 
 
 @pytest.fixture
-def runner(mocker: MockerFixture):
+def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[FunctionCallAgentRunner]:
     # Completely bypass BaseAgentRunner __init__ to avoid DB / Flask context
     mocker.patch(
         "core.agent.base_agent_runner.BaseAgentRunner.__init__",
@@ -131,7 +134,7 @@ def runner(mocker: MockerFixture):
     runner._current_thoughts = []
     runner.files = []
     runner.agent_callback = MagicMock()
-    runner.session = MagicMock()
+    runner.session = Session(sqlite_engine)
 
     runner._init_prompt_tools = MagicMock(return_value=({}, []))
     runner.create_agent_thought = MagicMock(return_value="thought1")
@@ -139,7 +142,10 @@ def runner(mocker: MockerFixture):
     runner.recalc_llm_max_tokens = MagicMock()
     runner.update_prompt_message_tool = MagicMock()
 
-    return runner
+    try:
+        yield runner
+    finally:
+        runner.session.close()
 
 
 # ==============================

api/tests/unit_tests/services/agent/test_skill_tool_inference_service.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 from unittest.mock import MagicMock, patch
 
 import pytest
+from sqlalchemy.orm import Session
 
 from services.agent.skill_tool_inference_service import (
     SkillToolInferenceError,
@@ -29,7 +30,8 @@ def _service(preview=_SKILL_MD_PREVIEW):
     return SkillToolInferenceService(drive_service=drive), drive
 
 
-def test_infer_returns_suggestions_with_inferred_from(monkeypatch):
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_infer_returns_suggestions_with_inferred_from(monkeypatch, sqlite_session: Session):
     service, drive = _service()
     raw = (
         '{"inferable": true, "reason": null, "cli_tools": [{"name": "ffmpeg",'
@@ -38,8 +40,12 @@ def test_infer_returns_suggestions_with_inferred_from(monkeypatch):
         ' "env_suggestions": [{"key": "OPENAI_API_KEY", "reason": "whisper call", "secret_likely": true}]}]}'
     )
     with patch.object(SkillToolInferenceService, "_invoke", staticmethod(lambda **kwargs: raw)):
-        session = MagicMock()
-        result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=session)
+        result = service.infer(
+            tenant_id="t-1",
+            agent_id="a-1",
+            slug="audio-transcribe",
+            session=sqlite_session,
+        )
 
     assert result["inferable"] is True
     tool = result["cli_tools"][0]
@@ -47,11 +53,13 @@ def test_infer_returns_suggestions_with_inferred_from(monkeypatch):
     assert tool["inferred_from"] == "audio-transcribe"
     assert tool["env_suggestions"] == [{"key": "OPENAI_API_KEY", "reason": "whisper call", "secret_likely": True}]
     drive.preview.assert_called_once_with(
-        tenant_id="t-1", agent_id="a-1", key="audio-transcribe/SKILL.md", session=session
+        tenant_id="t-1", agent_id="a-1", key="audio-transcribe/SKILL.md", session=sqlite_session
     )
+    assert not sqlite_session.in_transaction()
 
 
-def test_infer_threads_skill_md_into_the_prompt(monkeypatch):
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_infer_threads_skill_md_into_the_prompt(monkeypatch, sqlite_session: Session):
     service, _ = _service()
     captured: dict[str, str] = {}
 
@@ -60,21 +68,25 @@ def test_infer_threads_skill_md_into_the_prompt(monkeypatch):
         return '{"inferable": false, "cli_tools": [], "reason": "none"}'
 
     with patch.object(SkillToolInferenceService, "_invoke", staticmethod(fake_invoke)):
-        service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=MagicMock())
+        service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=sqlite_session)
 
     assert "Files inside the skill package" not in captured["prompt"]
     assert "ffmpeg" in captured["prompt"]  # SKILL.md body present
+    assert not sqlite_session.in_transaction()
 
 
-def test_infer_not_inferable_passes_reason_through(monkeypatch):
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_infer_not_inferable_passes_reason_through(monkeypatch, sqlite_session: Session):
     service, _ = _service()
     raw = '{"inferable": false, "cli_tools": [], "reason": "SKILL.md 未描述任何外部命令依赖"}'
     with patch.object(SkillToolInferenceService, "_invoke", staticmethod(lambda **kwargs: raw)):
-        result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=MagicMock())
+        result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=sqlite_session)
     assert result == {"inferable": False, "cli_tools": [], "reason": "SKILL.md 未描述任何外部命令依赖"}
+    assert not sqlite_session.in_transaction()
 
 
-def test_infer_retries_once_then_422(monkeypatch):
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_infer_retries_once_then_422(monkeypatch, sqlite_session: Session):
     service, _ = _service()
     calls: list[int] = []
 
@@ -84,37 +96,44 @@ def test_infer_retries_once_then_422(monkeypatch):
 
     with patch.object(SkillToolInferenceService, "_invoke", staticmethod(bad_invoke)):
         with pytest.raises(SkillToolInferenceError) as exc_info:
-            service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=MagicMock())
+            service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=sqlite_session)
 
     assert len(calls) == 2  # one retry
     assert exc_info.value.code == "inference_failed"
     assert exc_info.value.status_code == 422
+    assert not sqlite_session.in_transaction()
 
 
-def test_infer_repairs_slightly_malformed_json(monkeypatch):
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_infer_repairs_slightly_malformed_json(monkeypatch, sqlite_session: Session):
     service, _ = _service()
     raw = 'Here you go: {"inferable": true, "cli_tools": [], "reason": null,}'
     with patch.object(SkillToolInferenceService, "_invoke", staticmethod(lambda **kwargs: raw)):
-        result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=MagicMock())
+        result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=sqlite_session)
     assert result["inferable"] is True
+    assert not sqlite_session.in_transaction()
 
 
-def test_missing_skill_maps_to_404():
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_missing_skill_maps_to_404(sqlite_session: Session):
     drive = MagicMock()
     drive.preview.side_effect = AgentDriveError("drive_key_not_found", "nope", status_code=404)
     service = SkillToolInferenceService(drive_service=drive)
 
     with pytest.raises(SkillToolInferenceError) as exc_info:
-        service.infer(tenant_id="t-1", agent_id="a-1", slug="ghost", session=MagicMock())
+        service.infer(tenant_id="t-1", agent_id="a-1", slug="ghost", session=sqlite_session)
     assert exc_info.value.code == "skill_not_found"
     assert exc_info.value.status_code == 404
+    assert not sqlite_session.in_transaction()
 
 
-def test_binary_skill_md_maps_to_404():
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_binary_skill_md_maps_to_404(sqlite_session: Session):
     service, _ = _service(preview={"key": "x/SKILL.md", "size": 1, "truncated": False, "binary": True, "text": None})
     with pytest.raises(SkillToolInferenceError) as exc_info:
-        service.infer(tenant_id="t-1", agent_id="a-1", slug="x", session=MagicMock())
+        service.infer(tenant_id="t-1", agent_id="a-1", slug="x", session=sqlite_session)
     assert exc_info.value.code == "skill_not_found"
+    assert not sqlite_session.in_transaction()
 
 
 # ── real-path coverage: _invoke / passthrough ────────────────────────────────
@@ -157,11 +176,13 @@ def test_invoke_maps_model_failure_to_422_and_success_returns_text(monkeypatch):
     assert call["stream"] is False
 
 
-def test_load_skill_md_passes_through_non_missing_drive_errors():
+@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
+def test_load_skill_md_passes_through_non_missing_drive_errors(sqlite_session: Session):
     drive = MagicMock()
     drive.preview.side_effect = AgentDriveError("agent_not_found", "tenant mismatch", status_code=404)
     service = SkillToolInferenceService(drive_service=drive)
 
     with pytest.raises(SkillToolInferenceError) as exc_info:
-        service.infer(tenant_id="t-1", agent_id="a-1", slug="x", session=MagicMock())
+        service.infer(tenant_id="t-1", agent_id="a-1", slug="x", session=sqlite_session)
     assert exc_info.value.code == "agent_not_found"
+    assert not sqlite_session.in_transaction()