opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-01 15:09:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

test: migrate media and execution service sessions and ORM models to SQLite (#40552)

Browse Source
This commit is contained in:
Asuka Minato
2026-08-20 05:11:46 +00:00
committed by GitHub
parent 2155adff5a
commit e27eb51f8a
7 changed files with 408 additions and 221 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/tests/unit_tests/controllers/console/app/test_audio.py
+61 -34
View File
@@ -2,12 +2,12 @@ from __future__ import annotations
import io
from inspect import unwrap
from types import SimpleNamespace
from unittest.mock import patch
from uuid import UUID
import pytest
from flask import Flask
from sqlalchemy.orm import Session
from werkzeug.datastructures import FileStorage
from werkzeug.exceptions import Forbidden, InternalServerError
@@ -34,7 +34,7 @@ from controllers.console.app.error import (
)
from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
from graphon.model_runtime.errors.invoke import InvokeError
from models import AppMode
from models import Account, App, AppMode
from models.agent import AgentConfigDraftType
from models.agent_config_entities import AgentSoulConfig
from services.agent.composer_service import AgentComposerService
@@ -56,11 +56,30 @@ def _file_data():
return FileStorage(stream=io.BytesIO(b"audio"), filename="audio.wav", content_type="audio/wav")
def _app(*, app_id: str = "a1", tenant_id: str = "tenant-1") -> App:
return App(
id=app_id,
tenant_id=tenant_id,
name="Audio app",
description="",
mode=AppMode.CHAT,
enable_site=True,
enable_api=True,
max_active_requests=0,
)
def _account(account_id: str = "account-1") -> Account:
account = Account(name="Audio account", email=f"{account_id}@example.com")
account.id = account_id
return account
def test_console_audio_api_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: {"text": "ok"})
api = ChatMessageAudioApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1")
app_model = _app()
with app.test_request_context("/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()}):
response = handler(api, app_model=app_model)
@@ -72,9 +91,11 @@ def test_console_audio_api_accepts_published_agent_apps() -> None:
assert AppMode.AGENT in audio_module._CONSOLE_AUDIO_TRANSCRIPT_APP_MODES
def test_agent_console_audio_api_uses_agent_draft(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
def test_agent_console_audio_api_uses_agent_draft(
app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session
) -> None:
agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853")
app_model = SimpleNamespace(id="backing-app-1")
app_model = _app(app_id="backing-app-1")
agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": True}}})
calls: dict[str, object] = {}
@@ -100,8 +121,8 @@ def test_agent_console_audio_api_uses_agent_draft(app: Flask, monkeypatch: pytes
api = AgentChatMessageAudioApi()
handler = unwrap(api.post)
session = SimpleNamespace()
current_user = SimpleNamespace(id="account-1")
session = unbound_session
current_user = _account()
with app.test_request_context(
f"/console/api/agent/{agent_id}/audio-to-text",
method="POST",
@@ -140,13 +161,15 @@ def test_agent_console_audio_api_uses_agent_draft(app: Flask, monkeypatch: pytes
}
def test_agent_console_audio_api_defaults_to_normal_draft(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
def test_agent_console_audio_api_defaults_to_normal_draft(
app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session
) -> None:
agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853")
captured: dict[str, object] = {}
monkeypatch.setattr(
audio_module,
"resolve_agent_runtime_app_model",
lambda **_kwargs: SimpleNamespace(id="backing-app-1"),
lambda **_kwargs: _app(app_id="backing-app-1"),
)
def load_agent_soul_for_debug(**kwargs):
@@ -165,9 +188,9 @@ def test_agent_console_audio_api_defaults_to_normal_draft(app: Flask, monkeypatc
):
response = handler(
api,
session=SimpleNamespace(),
session=unbound_session,
current_tenant_id="tenant-1",
current_user=SimpleNamespace(id="account-1"),
current_user=_account(),
agent_id=agent_id,
)
@@ -175,9 +198,11 @@ def test_agent_console_audio_api_defaults_to_normal_draft(app: Flask, monkeypatc
assert captured["draft_type"] == AgentConfigDraftType.DRAFT
def test_agent_console_audio_api_checks_rbac_with_backing_app_id(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
def test_agent_console_audio_api_checks_rbac_with_backing_app_id(
app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session
) -> None:
agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853")
app_model = SimpleNamespace(id="backing-app-1")
app_model = _app(app_id="backing-app-1")
soul_loaded = False
monkeypatch.setattr(audio_module, "resolve_agent_runtime_app_model", lambda **_kwargs: app_model)
@@ -204,21 +229,23 @@ def test_agent_console_audio_api_checks_rbac_with_backing_app_id(app: Flask, mon
with pytest.raises(Forbidden):
handler(
api,
session=SimpleNamespace(),
session=unbound_session,
current_tenant_id="tenant-1",
current_user=SimpleNamespace(id="account-1"),
current_user=_account(),
agent_id=agent_id,
)
assert soul_loaded is False
def test_agent_console_audio_api_preserves_missing_build_draft_404(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
def test_agent_console_audio_api_preserves_missing_build_draft_404(
app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session
) -> None:
agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853")
monkeypatch.setattr(
audio_module,
"resolve_agent_runtime_app_model",
lambda **_kwargs: SimpleNamespace(id="backing-app-1"),
lambda **_kwargs: _app(app_id="backing-app-1"),
)
monkeypatch.setattr(
AgentComposerService,
@@ -236,9 +263,9 @@ def test_agent_console_audio_api_preserves_missing_build_draft_404(app: Flask, m
with pytest.raises(AgentVersionNotFoundError):
handler(
api,
session=SimpleNamespace(),
session=unbound_session,
current_tenant_id="tenant-1",
current_user=SimpleNamespace(id="account-1"),
current_user=_account(),
agent_id=agent_id,
)
@@ -262,7 +289,7 @@ def test_console_audio_api_error_mapping(app: Flask, monkeypatch: pytest.MonkeyP
monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: (_ for _ in ()).throw(exc))
api = ChatMessageAudioApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1")
app_model = _app()
with app.test_request_context("/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()}):
with pytest.raises(expected):
@@ -273,7 +300,7 @@ def test_console_audio_api_unhandled_error(app: Flask, monkeypatch: pytest.Monke
monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("boom")))
api = ChatMessageAudioApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1")
app_model = _app()
with app.test_request_context("/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()}):
with pytest.raises(InternalServerError):
@@ -285,7 +312,7 @@ def test_console_text_api_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -
api = ChatMessageTextApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1")
app_model = _app()
with app.test_request_context(
"/console/api/apps/app/text-to-audio",
@@ -300,7 +327,7 @@ def test_console_text_api_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -
def test_console_text_api_builds_message_ref(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None:
api = ChatMessageTextApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1")
app_model = _app(app_id="app-1")
calls = {}
def fake_transcript_tts(**kwargs):
@@ -315,7 +342,7 @@ def test_console_text_api_builds_message_ref(app: Flask, monkeypatch: pytest.Mon
method="POST",
json={"text": "hello", "message_id": "message-1"},
),
patch("controllers.console.app.audio.current_user", SimpleNamespace(id="account-1")),
patch("controllers.console.app.audio.current_user", _account()),
):
response = handler(api, TextToSpeechPayload(text="hello", message_id="message-1"), app_model=app_model)
@@ -328,7 +355,7 @@ def test_console_text_api_error_mapping(app: Flask, monkeypatch: pytest.MonkeyPa
api = ChatMessageTextApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1")
app_model = _app()
with app.test_request_context(
"/console/api/apps/app/text-to-audio",
@@ -345,7 +372,7 @@ def test_console_text_modes_success(app: Flask, monkeypatch: pytest.MonkeyPatch)
api = TextModesApi()
handler = unwrap(api.get)
app_model = SimpleNamespace(tenant_id="t1")
app_model = _app(tenant_id="t1")
with app.test_request_context("/console/api/apps/app/text-to-audio/voices?language=en", method="GET"):
response = handler(api, TextToSpeechVoiceQuery(language="en-US"), app_model=app_model)
@@ -362,7 +389,7 @@ def test_console_text_modes_language_error(app: Flask, monkeypatch: pytest.Monke
api = TextModesApi()
handler = unwrap(api.get)
app_model = SimpleNamespace(tenant_id="t1")
app_model = _app(tenant_id="t1")
with app.test_request_context("/console/api/apps/app/text-to-audio/voices?language=en", method="GET"):
with pytest.raises(AppUnavailableError):
@@ -376,7 +403,7 @@ def test_audio_to_text_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> N
response_payload = {"text": "hello"}
monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: response_payload)
app_model = SimpleNamespace(id="app-1")
app_model = _app(app_id="app-1")
data = {"file": (io.BytesIO(b"x"), "sample.wav")}
with app.test_request_context(
@@ -400,7 +427,7 @@ def test_audio_to_text_maps_audio_too_large(app: Flask, monkeypatch: pytest.Monk
lambda **_kwargs: (_ for _ in ()).throw(AudioTooLargeServiceError("too large")),
)
app_model = SimpleNamespace(id="app-1")
app_model = _app(app_id="app-1")
data = {"file": (io.BytesIO(b"x"), "sample.wav")}
with app.test_request_context(
@@ -419,7 +446,7 @@ def test_text_to_audio_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> N
monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: {"audio": "ok"})
app_model = SimpleNamespace(id="app-1")
app_model = _app(app_id="app-1")
with app.test_request_context(
"/console/api/apps/app-1/text-to-audio",
@@ -438,7 +465,7 @@ def test_text_to_audio_voices_success(app: Flask, monkeypatch: pytest.MonkeyPatc
expected_voices = [{"name": "Voice 1", "value": "voice-1"}]
monkeypatch.setattr(AudioService, "transcript_tts_voices", lambda **_kwargs: expected_voices)
app_model = SimpleNamespace(tenant_id="tenant-1")
app_model = _app()
with app.test_request_context(
"/console/api/apps/app-1/text-to-audio/voices",
@@ -456,7 +483,7 @@ def test_audio_to_text_with_invalid_file(app: Flask, monkeypatch: pytest.MonkeyP
monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: {"text": "test"})
app_model = SimpleNamespace(id="app-1")
app_model = _app(app_id="app-1")
data = {"file": (io.BytesIO(b"invalid"), "sample.xyz")}
with app.test_request_context(
@@ -476,7 +503,7 @@ def test_text_to_audio_with_language_param(app: Flask, monkeypatch: pytest.Monke
monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: {"audio": "test"})
app_model = SimpleNamespace(id="app-1")
app_model = _app(app_id="app-1")
with app.test_request_context(
"/console/api/apps/app-1/text-to-audio",
@@ -497,7 +524,7 @@ def test_text_to_audio_voices_with_language_filter(app: Flask, monkeypatch: pyte
lambda **_kwargs: [{"name": "Voice 1", "value": "voice-1"}],
)
app_model = SimpleNamespace(tenant_id="tenant-1")
app_model = _app()
with app.test_request_context(
"/console/api/apps/app-1/text-to-audio/voices?language=en-US",
api/tests/unit_tests/controllers/service_api/app/test_audio.py
+39 -13
View File
@@ -10,7 +10,6 @@ Tests coverage for:
import io
import uuid
from inspect import unwrap
from types import SimpleNamespace
from unittest.mock import Mock, patch
import pytest
@@ -34,6 +33,8 @@ from controllers.service_api.app.error import (
)
from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
from graphon.model_runtime.errors.invoke import InvokeError
from models.enums import EndUserType
from models.model import App, AppMode, EndUser
from services.app_ref_service import AppRef, MessageRef
from services.audio_service import AudioService
from services.errors.app_model_config import AppModelConfigBrokenError
@@ -50,6 +51,31 @@ def _file_data():
return FileStorage(stream=io.BytesIO(b"audio"), filename="audio.wav", content_type="audio/wav")
def _app(*, app_id: str = "a1", tenant_id: str = "tenant-1") -> App:
return App(
id=app_id,
tenant_id=tenant_id,
name="Audio app",
description="",
mode=AppMode.CHAT,
enable_site=True,
enable_api=True,
max_active_requests=0,
)
def _end_user(*, end_user_id: str = "u1", external_user_id: str | None = None) -> EndUser:
return EndUser(
id=end_user_id,
tenant_id="tenant-1",
app_id="a1",
type=EndUserType.SERVICE_API,
external_user_id=external_user_id,
name="Audio user",
session_id=f"session-{end_user_id}",
)
# ---------------------------------------------------------------------------
# Pydantic Model Tests
# ---------------------------------------------------------------------------
@@ -197,8 +223,8 @@ class TestAudioApi:
monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: {"text": "ok"})
api = AudioApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1")
end_user = SimpleNamespace(id="u1")
app_model = _app()
end_user = _end_user()
with app.test_request_context("/audio-to-text", method="POST", data={"file": _file_data()}):
response = handler(api, app_model=app_model, end_user=end_user)
@@ -224,8 +250,8 @@ class TestAudioApi:
monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: (_ for _ in ()).throw(exc))
api = AudioApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1")
end_user = SimpleNamespace(id="u1")
app_model = _app()
end_user = _end_user()
with app.test_request_context("/audio-to-text", method="POST", data={"file": _file_data()}):
with pytest.raises(expected):
@@ -237,8 +263,8 @@ class TestAudioApi:
)
api = AudioApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1")
end_user = SimpleNamespace(id="u1")
app_model = _app()
end_user = _end_user()
with app.test_request_context("/audio-to-text", method="POST", data={"file": _file_data()}):
with pytest.raises(InternalServerError):
@@ -251,8 +277,8 @@ class TestTextApi:
api = TextApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1")
end_user = SimpleNamespace(id="end-user-1", external_user_id="ext")
app_model = _app()
end_user = _end_user(end_user_id="end-user-1", external_user_id="ext")
with app.test_request_context(
"/text-to-audio",
@@ -274,8 +300,8 @@ class TestTextApi:
api = TextApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1", tenant_id="tenant-1")
end_user = SimpleNamespace(id="end-user-1", external_user_id="ext")
app_model = _app()
end_user = _end_user(end_user_id="end-user-1", external_user_id="ext")
with app.test_request_context(
"/text-to-audio",
@@ -294,8 +320,8 @@ class TestTextApi:
api = TextApi()
handler = unwrap(api.post)
app_model = SimpleNamespace(id="a1")
end_user = SimpleNamespace(id="end-user-1", external_user_id="ext")
app_model = _app()
end_user = _end_user(end_user_id="end-user-1", external_user_id="ext")
with app.test_request_context("/text-to-audio", method="POST", json={"text": "hello"}):
with pytest.raises(ProviderQuotaExceededError):
api/tests/unit_tests/services/test_api_token_service.py
+23 -11
View File
@@ -14,10 +14,29 @@ from werkzeug.exceptions import Unauthorized
import services.api_token_service as api_token_service_module
from models.engine import db
from models.enums import ApiTokenType
from models.model import ApiToken
from services.api_token_service import ApiTokenCache, CachedApiToken
def _api_token(
*,
token_id: str = "id-123",
app_id: str = "app-123",
tenant_id: str = "tenant-123",
token: str = "token-123",
) -> ApiToken:
"""Create a mapped API token for cache and single-flight behavior tests."""
return ApiToken(
id=token_id,
app_id=app_id,
tenant_id=tenant_id,
type=ApiTokenType.APP,
token=token,
last_used_at=None,
)
@pytest.fixture
def api_token_db() -> Iterator[Session]:
"""Provide the production database extension with an isolated SQLite token table."""
@@ -124,7 +143,7 @@ class TestFetchTokenWithSingleFlight:
def test_should_query_db_when_lock_acquired_and_cache_missed(self):
auth_token = "token-123"
scope = "app"
db_token = MagicMock()
db_token = _api_token()
lock = MagicMock()
lock.acquire.return_value = True
@@ -143,7 +162,7 @@ class TestFetchTokenWithSingleFlight:
def test_should_query_db_directly_when_lock_not_acquired(self):
auth_token = "token-123"
scope = "app"
db_token = MagicMock()
db_token = _api_token()
lock = MagicMock()
lock.acquire.return_value = False
@@ -184,7 +203,7 @@ class TestFetchTokenWithSingleFlight:
def test_should_fallback_to_db_query_when_lock_raises_exception(self):
auth_token = "token-123"
scope = "app"
db_token = MagicMock()
db_token = _api_token()
lock = MagicMock()
lock.acquire.side_effect = RuntimeError("redis lock error")
@@ -322,14 +341,7 @@ class TestApiTokenCacheCoreBranches:
@patch("services.api_token_service.redis_client")
def test_set_should_return_false_when_cache_write_raises_exception(self, mock_redis):
mock_redis.setex.side_effect = Exception("redis write failed")
api_token = MagicMock()
api_token.id = "id-123"
api_token.app_id = "app-123"
api_token.tenant_id = "tenant-123"
api_token.type = "app"
api_token.token = "token-123"
api_token.last_used_at = None
api_token.created_at = None
api_token = _api_token()
result = ApiTokenCache.set("token-123", "app", api_token)
assert result is False
api/tests/unit_tests/services/test_async_workflow_service.py
+58 -31
View File
@@ -10,9 +10,11 @@ from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session
import services.async_workflow_service as async_workflow_service_module
from models.enums import AppTriggerType, CreatorUserRole, WorkflowRunTriggeredFrom, WorkflowTriggerStatus
from models.model import App, AppMode
from models.account import Account
from models.enums import AppTriggerType, CreatorUserRole, EndUserType, WorkflowRunTriggeredFrom, WorkflowTriggerStatus
from models.model import App, AppMode, EndUser
from models.trigger import WorkflowTriggerLog
from models.workflow import Workflow, WorkflowType
from services.async_workflow_service import AsyncWorkflowService
from services.errors.app import QuotaExceededError, WorkflowNotFoundError
from services.workflow.entities import AsyncTriggerResponse, TriggerData
@@ -94,6 +96,41 @@ class AsyncWorkflowServiceTestDataFactory:
trigger_log.created_at = created_at
return trigger_log
@staticmethod
def create_workflow(
*, workflow_id: str = "workflow-123", app_id: str = "app-123", tenant_id: str = "tenant-123"
) -> Workflow:
"""Create a mapped workflow for service-return and trigger-log tests."""
return Workflow(
id=workflow_id,
tenant_id=tenant_id,
app_id=app_id,
type=WorkflowType.WORKFLOW,
version="1",
graph="{}",
_features="{}",
created_by="account-123",
)
@staticmethod
def create_account(account_id: str = "account-123") -> Account:
"""Create a mapped account for role-discrimination tests."""
account = Account(name="Account", email=f"{account_id}@example.com")
account.id = account_id
return account
@staticmethod
def create_end_user(end_user_id: str = "end-user-123") -> EndUser:
"""Create a mapped end user for role-discrimination and retry tests."""
return EndUser(
id=end_user_id,
tenant_id="tenant-123",
app_id="app-123",
type=EndUserType.BROWSER,
name="End User",
session_id=f"session-{end_user_id}",
)
@pytest.mark.usefixtures("sqlite_session")
@pytest.mark.parametrize("sqlite_session", [(App, WorkflowTriggerLog)], indirect=True)
@@ -163,8 +200,7 @@ class TestAsyncWorkflowService:
sqlite_session.add(AsyncWorkflowServiceTestDataFactory.create_app())
sqlite_session.commit()
trigger_data = AsyncWorkflowServiceTestDataFactory.create_trigger_data()
workflow = MagicMock()
workflow.id = "workflow-123"
workflow = AsyncWorkflowServiceTestDataFactory.create_workflow()
mocks = async_workflow_trigger_mocks
mocks["dispatcher"].get_queue_name.return_value = queue_name
@@ -179,17 +215,12 @@ class TestAsyncWorkflowService:
quota_charge_mock = MagicMock()
mocks["quota_service"].reserve.return_value = quota_charge_mock
class DummyAccount:
def __init__(self, user_id: str):
self.id = user_id
user = AsyncWorkflowServiceTestDataFactory.create_account()
with patch.object(async_workflow_service_module, "Account", DummyAccount):
user = DummyAccount("account-123")
# Act
result = AsyncWorkflowService.trigger_workflow_async(
session=sqlite_session, user=user, trigger_data=trigger_data
)
# Act
result = AsyncWorkflowService.trigger_workflow_async(
session=sqlite_session, user=user, trigger_data=trigger_data
)
# Assert
assert isinstance(result, AsyncTriggerResponse)
@@ -233,8 +264,7 @@ class TestAsyncWorkflowService:
sqlite_session.add(AsyncWorkflowServiceTestDataFactory.create_app())
sqlite_session.commit()
trigger_data = AsyncWorkflowServiceTestDataFactory.create_trigger_data()
workflow = MagicMock()
workflow.id = "workflow-123"
workflow = AsyncWorkflowServiceTestDataFactory.create_workflow()
mocks = async_workflow_trigger_mocks
mocks["dispatcher"].get_queue_name.return_value = QueuePriority.SANDBOX
@@ -243,7 +273,7 @@ class TestAsyncWorkflowService:
task_result = MagicMock(id="task-123")
mocks["sandbox_task"].delay.return_value = task_result
user = SimpleNamespace(id="end-user-123")
user = AsyncWorkflowServiceTestDataFactory.create_end_user()
# Act
response = AsyncWorkflowService.trigger_workflow_async(
@@ -269,7 +299,7 @@ class TestAsyncWorkflowService:
with pytest.raises(WorkflowNotFoundError, match="App not found: missing-app"):
AsyncWorkflowService.trigger_workflow_async(
session=sqlite_session,
user=SimpleNamespace(id="user-123"),
user=AsyncWorkflowServiceTestDataFactory.create_end_user("user-123"),
trigger_data=trigger_data,
)
@@ -284,8 +314,7 @@ class TestAsyncWorkflowService:
sqlite_session.add(AsyncWorkflowServiceTestDataFactory.create_app())
sqlite_session.commit()
trigger_data = AsyncWorkflowServiceTestDataFactory.create_trigger_data()
workflow = MagicMock()
workflow.id = "workflow-123"
workflow = AsyncWorkflowServiceTestDataFactory.create_workflow()
mocks = async_workflow_trigger_mocks
mocks["dispatcher"].get_queue_name.return_value = QueuePriority.TEAM
@@ -301,7 +330,7 @@ class TestAsyncWorkflowService:
with pytest.raises(QuotaExceededError) as exc_info:
AsyncWorkflowService.trigger_workflow_async(
session=sqlite_session,
user=SimpleNamespace(id="user-123"),
user=AsyncWorkflowServiceTestDataFactory.create_end_user("user-123"),
trigger_data=trigger_data,
)
@@ -328,7 +357,7 @@ class TestAsyncWorkflowService:
with pytest.raises(ValueError, match="Trigger log not found: missing-log"):
AsyncWorkflowService.reinvoke_trigger(
session=sqlite_session,
user=SimpleNamespace(id="user-123"),
user=AsyncWorkflowServiceTestDataFactory.create_end_user("user-123"),
workflow_trigger_log_id="missing-log",
)
@@ -354,7 +383,7 @@ class TestAsyncWorkflowService:
return_value=expected_response,
) as mock_trigger_workflow_async,
):
user = SimpleNamespace(id="user-123")
user = AsyncWorkflowServiceTestDataFactory.create_end_user("user-123")
# Act
response = AsyncWorkflowService.reinvoke_trigger(
@@ -497,8 +526,8 @@ class TestAsyncWorkflowServiceGetWorkflow:
"""Test _get_workflow returns published workflow by id when provided."""
# Arrange
workflow_service = MagicMock()
app_model = MagicMock()
workflow = MagicMock()
app_model = AsyncWorkflowServiceTestDataFactory.create_app()
workflow = AsyncWorkflowServiceTestDataFactory.create_workflow()
workflow_service.get_published_workflow_by_id.return_value = workflow
# Act
@@ -517,7 +546,7 @@ class TestAsyncWorkflowServiceGetWorkflow:
"""Test _get_workflow raises WorkflowNotFoundError for unknown workflow id."""
# Arrange
workflow_service = MagicMock()
app_model = MagicMock()
app_model = AsyncWorkflowServiceTestDataFactory.create_app()
workflow_service.get_published_workflow_by_id.return_value = None
# Act / Assert
@@ -530,9 +559,8 @@ class TestAsyncWorkflowServiceGetWorkflow:
"""Test _get_workflow returns default published workflow when no id is provided."""
# Arrange
workflow_service = MagicMock()
app_model = MagicMock()
app_model.id = "app-123"
workflow = MagicMock()
app_model = AsyncWorkflowServiceTestDataFactory.create_app()
workflow = AsyncWorkflowServiceTestDataFactory.create_workflow()
workflow_service.get_published_workflow.return_value = workflow
# Act
@@ -547,8 +575,7 @@ class TestAsyncWorkflowServiceGetWorkflow:
"""Test _get_workflow raises WorkflowNotFoundError when app has no published workflow."""
# Arrange
workflow_service = MagicMock()
app_model = MagicMock()
app_model.id = "app-123"
app_model = AsyncWorkflowServiceTestDataFactory.create_app()
workflow_service.get_published_workflow.return_value = None
# Act / Assert
api/tests/unit_tests/services/test_audio_service.py
+59 -33
View File
@@ -53,9 +53,11 @@ Tests available voice retrieval:
- text_to_speech: Enables TTS functionality
"""
import json
from decimal import Decimal
from typing import Any
from unittest.mock import MagicMock, Mock, create_autospec, patch
from unittest.mock import MagicMock, Mock, patch
from uuid import uuid4
import pytest
from sqlalchemy.orm import Session
@@ -64,7 +66,7 @@ from werkzeug.datastructures import FileStorage
from models.agent_config_entities import AgentSoulConfig
from models.enums import ConversationFromSource, MessageStatus
from models.model import App, AppMode, AppModelConfig, Message
from models.workflow import Workflow
from models.workflow import Workflow, WorkflowType
from services.app_ref_service import AppRef, MessageRef
from services.audio_service import AudioService
from services.errors.audio import (
@@ -113,15 +115,18 @@ class AudioServiceTestDataFactory:
audio-related operations.
"""
@staticmethod
def __init__(self, session: Session) -> None:
self.session = session
def create_app_mock(
app_id: str = "app-123",
self,
app_id: str = APP_ID,
mode: AppMode = AppMode.CHAT,
tenant_id: str = "tenant-123",
tenant_id: str = TENANT_ID,
**kwargs,
) -> Mock:
) -> App:
"""
Create a mock App object.
Create and persist an App model.
Args:
app_id: Unique identifier for the app
@@ -130,46 +135,65 @@ class AudioServiceTestDataFactory:
**kwargs: Additional attributes to set on the mock
Returns:
Mock App object with specified attributes
Persisted App model with specified attributes
"""
app = create_autospec(App, instance=True)
app.id = app_id
app.mode = mode
app.tenant_id = tenant_id
app.workflow = kwargs.get("workflow")
app.app_model_config = kwargs.get("app_model_config")
app.workflow_with_session.return_value = app.workflow
app.app_model_config_with_session.return_value = app.app_model_config
workflow = kwargs.pop("workflow", None)
app_model_config = kwargs.pop("app_model_config", None)
app = App(
id=app_id,
tenant_id=tenant_id,
name="Audio test app",
description="",
mode=mode,
icon_type=None,
icon=None,
icon_background=None,
enable_site=False,
enable_api=False,
workflow_id=workflow.id if workflow else None,
app_model_config_id=app_model_config.id if app_model_config else None,
)
for key, value in kwargs.items():
setattr(app, key, value)
self.session.add(app)
self.session.commit()
return app
@staticmethod
def create_workflow_mock(features_dict: dict[str, Any] | None = None, **kwargs) -> Mock:
def create_workflow_mock(self, features_dict: dict[str, Any] | None = None, **kwargs) -> Workflow:
"""
Create a mock Workflow object.
Create and persist a Workflow model.
Args:
features_dict: Dictionary of workflow features
**kwargs: Additional attributes to set on the mock
Returns:
Mock Workflow object with specified attributes
Persisted Workflow model with specified attributes
"""
workflow = create_autospec(Workflow, instance=True)
workflow.features_dict = features_dict or {}
workflow = Workflow(
id=kwargs.pop("id", str(uuid4())),
tenant_id=kwargs.pop("tenant_id", TENANT_ID),
app_id=kwargs.pop("app_id", APP_ID),
type=kwargs.pop("type", WorkflowType.CHAT),
version=kwargs.pop("version", Workflow.VERSION_DRAFT),
graph=kwargs.pop("graph", "{}"),
_features=json.dumps(features_dict or {}),
created_by=kwargs.pop("created_by", ACCOUNT_ID),
)
for key, value in kwargs.items():
setattr(workflow, key, value)
self.session.add(workflow)
self.session.commit()
return workflow
@staticmethod
def create_app_model_config_mock(
self,
speech_to_text_dict: dict[str, Any] | None = None,
text_to_speech_dict: dict[str, Any] | None = None,
**kwargs,
) -> Mock:
) -> AppModelConfig:
"""
Create a mock AppModelConfig object.
Create and persist an AppModelConfig model.
Args:
speech_to_text_dict: Speech-to-text configuration
@@ -177,13 +201,17 @@ class AudioServiceTestDataFactory:
**kwargs: Additional attributes to set on the mock
Returns:
Mock AppModelConfig object with specified attributes
Persisted AppModelConfig model with specified attributes
"""
config = create_autospec(AppModelConfig, instance=True)
config.speech_to_text_dict = speech_to_text_dict or {"enabled": False}
config.text_to_speech_dict = text_to_speech_dict or {"enabled": False}
config = AppModelConfig(
app_id=kwargs.pop("app_id", APP_ID),
speech_to_text=json.dumps(speech_to_text_dict or {"enabled": False}),
text_to_speech=json.dumps(text_to_speech_dict or {"enabled": False}),
)
for key, value in kwargs.items():
setattr(config, key, value)
self.session.add(config)
self.session.commit()
return config
@staticmethod
@@ -216,9 +244,9 @@ class AudioServiceTestDataFactory:
@pytest.fixture
def factory():
def factory(sqlite_session: Session) -> AudioServiceTestDataFactory:
"""Provide the test data factory to all tests."""
return AudioServiceTestDataFactory
return AudioServiceTestDataFactory(sqlite_session)
class TestAudioServiceASR:
@@ -385,7 +413,6 @@ class TestAudioServiceASR:
self, mock_model_manager_class, factory: AudioServiceTestDataFactory
):
app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True})
app_model_config.to_dict.return_value = {"speech_to_text": {"enabled": True}}
app = factory.create_app_mock(mode=AppMode.AGENT, app_model_config=app_model_config)
file = factory.create_file_storage_mock()
mock_model_instance = MagicMock()
@@ -403,7 +430,6 @@ class TestAudioServiceASR:
def test_transcript_agent_asr_soul_disabled_overrides_legacy_feature(self, factory: AudioServiceTestDataFactory):
app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True})
app_model_config.to_dict.return_value = {"speech_to_text": {"enabled": True}}
app = factory.create_app_mock(mode=AppMode.AGENT, app_model_config=app_model_config)
file = factory.create_file_storage_mock()
agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": False}}})
api/tests/unit_tests/services/test_message_service.py
+17 -5
View File
@@ -38,6 +38,7 @@ from models.model import (
Message,
MessageFeedback,
)
from models.workflow import Workflow, WorkflowType
from repositories.sqlalchemy_execution_extra_content_repository import SQLAlchemyExecutionExtraContentRepository
from services.errors.message import (
FirstMessageNotExistsError,
@@ -111,6 +112,19 @@ class MessageServiceTestDataFactory:
account.id = user_id
return account
@staticmethod
def create_workflow(*, features: dict[str, object] | None = None) -> Workflow:
return Workflow(
id="workflow-123",
tenant_id="tenant-123",
app_id="app-123",
type=WorkflowType.CHAT,
version="1",
graph="{}",
_features=json.dumps(features or {}),
created_by="account-123",
)
@staticmethod
def create_conversation(
conversation_id: str = "conv-001",
@@ -712,8 +726,7 @@ class TestMessageServiceSuggestedQuestions:
monkeypatch: pytest.MonkeyPatch,
conversation: Conversation,
) -> tuple[MagicMock, MagicMock, MagicMock]:
message = MagicMock()
message.conversation_id = conversation.id
message = MessageServiceTestDataFactory.create_message(message_id="msg-123", conversation_id=conversation.id)
monkeypatch.setattr(service_module.MessageService, "get_message", MagicMock(return_value=message))
monkeypatch.setattr(
service_module.ConversationService, "get_conversation", MagicMock(return_value=conversation)
@@ -747,8 +760,7 @@ class TestMessageServiceSuggestedQuestions:
) -> None:
conversation = factory.create_conversation()
_, _, llm_generator = self._chat_boundaries(monkeypatch, conversation)
workflow = MagicMock()
workflow.features_dict = {"suggested_questions_after_answer": {"enabled": True}}
workflow = factory.create_workflow(features={"suggested_questions_after_answer": {"enabled": True}})
workflow_service = MagicMock()
workflow_service.return_value.get_published_workflow.return_value = workflow
monkeypatch.setattr(service_module, "WorkflowService", workflow_service)
@@ -1103,7 +1115,7 @@ class TestMessageServiceSuggestedQuestions:
) -> None:
conversation = factory.create_conversation()
self._chat_boundaries(monkeypatch, conversation)
workflow = MagicMock()
workflow = factory.create_workflow()
workflow_service = MagicMock()
workflow_service.return_value.get_published_workflow.return_value = workflow
monkeypatch.setattr(service_module, "WorkflowService", workflow_service)
api/tests/unit_tests/services/test_vector_service.py
+151 -94
View File
@@ -16,7 +16,7 @@ import services.vector_service as vector_service_module
from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from extensions.storage.storage_type import StorageType
from models import UploadFile
from models.dataset import ChildChunk, DatasetProcessRule, SegmentAttachmentBinding
from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment, SegmentAttachmentBinding
from models.dataset import Document as DatasetDocument
from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, ProcessRuleMode
from services.vector_service import VectorService
@@ -42,16 +42,18 @@ def _make_dataset(
is_multimodal: bool = False,
embedding_model_provider: str | None = "openai",
embedding_model: str = "text-embedding",
) -> MagicMock:
dataset = MagicMock(name="dataset")
dataset.id = dataset_id
dataset.tenant_id = tenant_id
dataset.doc_form = doc_form
dataset.indexing_technique = indexing_technique
dataset.is_multimodal = is_multimodal
dataset.embedding_model_provider = embedding_model_provider
dataset.embedding_model = embedding_model
dataset.get_doc_form.return_value = doc_form
) -> Dataset:
dataset = Dataset(
id=dataset_id,
tenant_id=tenant_id,
name="Dataset",
created_by="account-1",
indexing_technique=indexing_technique,
chunk_structure=doc_form,
is_multimodal=is_multimodal,
embedding_model_provider=embedding_model_provider,
embedding_model=embedding_model,
)
return dataset
@@ -64,24 +66,49 @@ def _make_segment(
content: str = "hello",
index_node_id: str = "node-1",
index_node_hash: str = "hash-1",
session: Session | None = None,
attachments: list[dict[str, str]] | None = None,
) -> MagicMock:
segment = MagicMock(name="segment")
) -> DocumentSegment:
segment = DocumentSegment(
tenant_id=tenant_id,
dataset_id=dataset_id,
document_id=document_id,
position=1,
content=content,
word_count=len(content),
tokens=len(content),
created_by="account-1",
index_node_id=index_node_id,
index_node_hash=index_node_hash,
)
segment.id = segment_id
segment.tenant_id = tenant_id
segment.dataset_id = dataset_id
segment.document_id = document_id
segment.content = content
segment.index_node_id = index_node_id
segment.index_node_hash = index_node_hash
segment.attachments = attachments or []
segment.get_attachments.return_value = attachments or []
if attachments:
assert session is not None
for attachment in attachments:
upload_file = _upload_file(
file_id=attachment["id"],
name=attachment.get("name", f"{attachment['id']}.png"),
tenant_id=tenant_id,
)
session.add_all(
[
upload_file,
SegmentAttachmentBinding(
tenant_id=tenant_id,
dataset_id=dataset_id,
document_id=document_id,
segment_id=segment_id,
attachment_id=upload_file.id,
),
]
)
session.flush()
return segment
def _upload_file(*, file_id: str = "file-1", name: str = "img.png") -> UploadFile:
def _upload_file(*, file_id: str = "file-1", name: str = "img.png", tenant_id: str = "tenant-1") -> UploadFile:
upload_file = UploadFile(
tenant_id="tenant-1",
tenant_id=tenant_id,
storage_type=StorageType.LOCAL,
key=f"uploads/{file_id}",
name=name,
@@ -97,6 +124,30 @@ def _upload_file(*, file_id: str = "file-1", name: str = "img.png") -> UploadFil
return upload_file
def _make_child_chunk(
*,
index_node_id: str,
content: str = "child",
index_node_hash: str = "hash",
tenant_id: str = "tenant-1",
dataset_id: str = "dataset-1",
document_id: str = "doc-1",
segment_id: str = "seg-1",
) -> ChildChunk:
return ChildChunk(
tenant_id=tenant_id,
dataset_id=dataset_id,
document_id=document_id,
segment_id=segment_id,
position=1,
content=content,
word_count=len(content),
created_by="account-1",
index_node_id=index_node_id,
index_node_hash=index_node_hash,
)
def test_create_segments_vector_regular_indexing_loads_documents_and_keywords(
monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
) -> None:
@@ -126,10 +177,11 @@ def test_create_segments_vector_regular_indexing_loads_multimodal_documents(
) -> None:
dataset = _make_dataset(is_multimodal=True)
segment = _make_segment(
session=sqlite_session,
attachments=[
{"id": "img-1", "name": "a.png"},
{"id": "img-2", "name": "b.png"},
]
],
)
index_processor = MagicMock(name="index_processor")
@@ -152,7 +204,7 @@ def test_create_segments_vector_regular_indexing_loads_multimodal_documents(
assert second_args[1] == []
assert len(second_args[2]) == 2
assert second_kwargs["with_keywords"] is False
segment.get_attachments.assert_called_once_with(session=sqlite_session)
assert {document.page_content for document in second_args[2]} == {"a.png", "b.png"}
def test_create_segments_vector_with_no_segments_does_not_load(
@@ -171,7 +223,7 @@ def test_create_segments_vector_with_no_segments_does_not_load(
def _persist_parent_child_rows(
session: Session,
*,
segment: MagicMock,
segment: DocumentSegment,
include_document: bool = True,
include_rule: bool = True,
) -> tuple[DatasetDocument | None, DatasetProcessRule | None]:
@@ -415,13 +467,24 @@ def test_generate_child_chunks_regenerate_cleans_then_saves_children(
dataset = _make_dataset(doc_form=IndexStructureType.PARAGRAPH_INDEX, tenant_id="tenant-1", dataset_id="dataset-1")
segment = _make_segment(segment_id="seg-1")
dataset_document = MagicMock()
dataset_document.id = segment.document_id
dataset_document.doc_language = "en"
dataset_document.created_by = "user-1"
processing_rule = MagicMock()
processing_rule.to_dict.return_value = {"rules": {}}
dataset_document = DatasetDocument(
id=segment.document_id,
tenant_id=segment.tenant_id,
dataset_id=segment.dataset_id,
position=1,
data_source_type=DataSourceType.UPLOAD_FILE,
batch="batch-1",
name="Document",
created_from=DocumentCreatedFrom.API,
created_by="user-1",
doc_language="en",
)
processing_rule = DatasetProcessRule(
dataset_id=segment.dataset_id,
mode=ProcessRuleMode.HIERARCHICAL,
rules="{}",
created_by="user-1",
)
child1 = _ChildDocStub(page_content="c1", metadata={"doc_id": "c1-id", "doc_hash": "c1-h"})
child2 = _ChildDocStub(page_content="c2", metadata={"doc_id": "c2-id", "doc_hash": "c2-h"})
@@ -456,12 +519,24 @@ def test_generate_child_chunks_flushes_even_when_no_children(
) -> None:
dataset = _make_dataset(doc_form=IndexStructureType.PARAGRAPH_INDEX)
segment = _make_segment()
dataset_document = MagicMock()
dataset_document.doc_language = "en"
dataset_document.created_by = "user-1"
processing_rule = MagicMock()
processing_rule.to_dict.return_value = {"rules": {}}
dataset_document = DatasetDocument(
id=segment.document_id,
tenant_id=segment.tenant_id,
dataset_id=segment.dataset_id,
position=1,
data_source_type=DataSourceType.UPLOAD_FILE,
batch="batch-1",
name="Document",
created_from=DocumentCreatedFrom.API,
created_by="user-1",
doc_language="en",
)
processing_rule = DatasetProcessRule(
dataset_id=segment.dataset_id,
mode=ProcessRuleMode.HIERARCHICAL,
rules="{}",
created_by="user-1",
)
index_processor = MagicMock()
index_processor.transform.return_value = [_ParentDocStub(children=[])]
@@ -487,12 +562,7 @@ def test_create_child_chunk_vector_high_quality_adds_texts(
monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
) -> None:
dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
child_chunk = MagicMock()
child_chunk.content = "child"
child_chunk.index_node_id = "id"
child_chunk.index_node_hash = "h"
child_chunk.document_id = "doc-1"
child_chunk.dataset_id = "dataset-1"
child_chunk = _make_child_chunk(index_node_id="id", index_node_hash="h")
vector_instance = MagicMock()
vector_cls = MagicMock(return_value=vector_instance)
@@ -508,12 +578,7 @@ def test_create_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch,
vector_cls = MagicMock()
monkeypatch.setattr(vector_service_module, "Vector", vector_cls)
child_chunk = MagicMock()
child_chunk.content = "child"
child_chunk.index_node_id = "id"
child_chunk.index_node_hash = "h"
child_chunk.document_id = "doc-1"
child_chunk.dataset_id = "dataset-1"
child_chunk = _make_child_chunk(index_node_id="id", index_node_hash="h")
VectorService.create_child_chunk_vector(child_chunk, dataset, session=sqlite_session)
vector_cls.assert_not_called()
@@ -524,22 +589,13 @@ def test_update_child_chunk_vector_high_quality_updates_vector(
) -> None:
dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY)
new_chunk = MagicMock()
new_chunk.content = "n"
new_chunk.index_node_id = "nid"
new_chunk.index_node_hash = "nh"
new_chunk.document_id = "d"
new_chunk.dataset_id = "ds"
upd_chunk = MagicMock()
upd_chunk.content = "u"
upd_chunk.index_node_id = "uid"
upd_chunk.index_node_hash = "uh"
upd_chunk.document_id = "d"
upd_chunk.dataset_id = "ds"
del_chunk = MagicMock()
del_chunk.index_node_id = "did"
new_chunk = _make_child_chunk(
content="n", index_node_id="nid", index_node_hash="nh", document_id="d", dataset_id="ds"
)
upd_chunk = _make_child_chunk(
content="u", index_node_id="uid", index_node_hash="uh", document_id="d", dataset_id="ds"
)
del_chunk = _make_child_chunk(index_node_id="did")
vector_instance = MagicMock()
vector_cls = MagicMock(return_value=vector_instance)
@@ -564,8 +620,7 @@ def test_update_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch,
def test_delete_child_chunk_vector_deletes_by_id(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None:
dataset = _make_dataset()
child_chunk = MagicMock()
child_chunk.index_node_id = "cid"
child_chunk = _make_child_chunk(index_node_id="cid")
vector_instance = MagicMock()
vector_cls = MagicMock(return_value=vector_instance)
@@ -585,7 +640,7 @@ def test_update_multimodel_vector_returns_when_not_high_quality(
monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
) -> None:
dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY, is_multimodal=True)
segment = _make_segment(tenant_id="t", attachments=[{"id": "a"}])
segment = _make_segment(tenant_id="t")
vector_cls = MagicMock()
monkeypatch.setattr(vector_service_module, "Vector", vector_cls)
@@ -601,7 +656,7 @@ def test_update_multimodel_vector_returns_when_no_actual_change(
monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
) -> None:
dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True)
segment = _make_segment(tenant_id="t", attachments=[{"id": "a"}, {"id": "b"}])
segment = _make_segment(tenant_id="t", session=sqlite_session, attachments=[{"id": "a"}, {"id": "b"}])
vector_cls = MagicMock()
monkeypatch.setattr(vector_service_module, "Vector", vector_cls)
@@ -610,7 +665,7 @@ def test_update_multimodel_vector_returns_when_no_actual_change(
session=sqlite_session, segment=segment, attachment_ids=["b", "a"], dataset=dataset
)
vector_cls.assert_not_called()
assert not sqlite_session.in_transaction()
assert sqlite_session.in_transaction()
def test_update_multimodel_vector_deletes_bindings_and_commits_on_empty_new_ids(
@@ -618,23 +673,14 @@ def test_update_multimodel_vector_deletes_bindings_and_commits_on_empty_new_ids(
sqlite_session: Session,
) -> None:
dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True)
segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}, {"id": "old-2"}])
segment = _make_segment(
tenant_id="tenant-1",
session=sqlite_session,
attachments=[{"id": "old-1"}, {"id": "old-2"}],
)
vector_instance = MagicMock(name="vector_instance")
vector_cls = MagicMock(return_value=vector_instance)
sqlite_session.add_all(
[
SegmentAttachmentBinding(
tenant_id="tenant-1",
dataset_id="dataset-1",
document_id="doc-1",
segment_id="seg-1",
attachment_id=attachment_id,
)
for attachment_id in ("old-1", "old-2")
]
)
sqlite_session.flush()
monkeypatch.setattr(vector_service_module, "Vector", vector_cls)
@@ -650,7 +696,7 @@ def test_update_multimodel_vector_flushes_when_no_upload_files_found(
monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
) -> None:
dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True)
segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}])
segment = _make_segment(tenant_id="tenant-1", session=sqlite_session, attachments=[{"id": "old-1"}])
vector_instance = MagicMock()
monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance))
@@ -668,7 +714,12 @@ def test_update_multimodel_vector_adds_bindings_and_vectors_and_skips_missing_up
sqlite_session: Session,
) -> None:
dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True)
segment = _make_segment(segment_id="seg-1", tenant_id="tenant-1", attachments=[{"id": "old-1"}])
segment = _make_segment(
segment_id="seg-1",
tenant_id="tenant-1",
session=sqlite_session,
attachments=[{"id": "old-1"}],
)
vector_instance = MagicMock()
monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance))
@@ -696,7 +747,7 @@ def test_update_multimodel_vector_updates_bindings_without_multimodal_vector_ops
sqlite_session: Session,
) -> None:
dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=False)
segment = _make_segment(tenant_id="tenant-1", attachments=[{"id": "old-1"}])
segment = _make_segment(tenant_id="tenant-1", session=sqlite_session, attachments=[{"id": "old-1"}])
vector_instance = MagicMock()
monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance))
@@ -719,7 +770,12 @@ def test_update_multimodel_vector_rolls_back_and_reraises_on_error(
sqlite_session: Session,
) -> None:
dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True)
segment = _make_segment(segment_id="seg-1", tenant_id="tenant-1", attachments=[{"id": "old-1"}])
segment = _make_segment(
segment_id="seg-1",
tenant_id="tenant-1",
session=sqlite_session,
attachments=[{"id": "old-1"}],
)
vector_instance = MagicMock()
monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance))
@@ -730,10 +786,11 @@ def test_update_multimodel_vector_rolls_back_and_reraises_on_error(
monkeypatch.setattr(sqlite_session, "flush", MagicMock(side_effect=RuntimeError("boom")))
with caplog.at_level(logging.ERROR, logger="services.vector_service"):
with pytest.raises(RuntimeError, match="boom"):
VectorService.update_multimodel_vector(
session=sqlite_session, segment=segment, attachment_ids=["file-1"], dataset=dataset
)
with sqlite_session.no_autoflush:
with pytest.raises(RuntimeError, match="boom"):
VectorService.update_multimodel_vector(
session=sqlite_session, segment=segment, attachment_ids=["file-1"], dataset=dataset
)
assert any(r.levelno >= logging.ERROR for r in caplog.records)
assert rollback_events == ["rollback"]