diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index ea6368075cb..9599be7aeb8 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/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: diff --git a/api/tests/unit_tests/controllers/console/test_wraps.py b/api/tests/unit_tests/controllers/console/test_wraps.py index 07ac57b1ede..995877a2507 100644 --- a/api/tests/unit_tests/controllers/console/test_wraps.py +++ b/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: diff --git a/api/tests/unit_tests/core/agent/test_cot_agent_runner.py b/api/tests/unit_tests/core/agent/test_cot_agent_runner.py index 43987bd65d0..d3cdc6ab292 100644 --- a/api/tests/unit_tests/core/agent/test_cot_agent_runner.py +++ b/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: diff --git a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py index d70605c96ef..9ce87271bf7 100644 --- a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py +++ b/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() # ============================== diff --git a/api/tests/unit_tests/services/agent/test_skill_tool_inference_service.py b/api/tests/unit_tests/services/agent/test_skill_tool_inference_service.py index cfb32d63a92..dd61ba79c19 100644 --- a/api/tests/unit_tests/services/agent/test_skill_tool_inference_service.py +++ b/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()