api/tests/unit_tests/conftest.py
Rama main · réplica de lectura del repositorio original.
from __future__ import annotations
import os
import shutil
from collections.abc import Callable, Iterator
from pathlib import Path
from tempfile import TemporaryDirectory
from threading import Lock
from typing import TYPE_CHECKING
from unittest.mock import MagicMock, create_autospec, patch
from weakref import WeakKeyDictionary
# Ensure OpenDAL fs writes to tmp to avoid polluting workspace
os.environ.setdefault("OPENDAL_SCHEME", "fs")
os.environ.setdefault("OPENDAL_FS_ROOT", "/tmp/dify-storage")
os.environ.setdefault("STORAGE_TYPE", "opendal")
import pytest
from dify_trace_sdk.testing import RecordingEgress
from flask import Flask
from redis import Redis
from sqlalchemy import create_engine, event
from sqlalchemy.engine import URL, Engine
from sqlalchemy.orm import Session, sessionmaker
import core.db.session_factory as session_factory_module
from extensions import ext_redis
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.base import TypeBase
from tests.helpers.unit.config_override import apply_config_overrides
from tests.helpers.unit.rbac_fakes import RBACDomain, build_rbac_domain
from tests.helpers.unit.services.tracing_ports import deployment_provider_configs
if TYPE_CHECKING:
from extensions.ext_application_services import ApplicationServices
from tests.helpers.unit.account_domain import AccountDomain
# Getting the absolute path of the current file's directory
ABS_PATH = os.path.dirname(os.path.abspath(__file__))
# Getting the absolute path of the project's root directory
PROJECT_DIR = os.path.abspath(os.path.join(ABS_PATH, os.pardir, os.pardir))
CACHED_APP = Flask(__name__)
pytest_plugins = (
"tests.helpers.unit.audio_runtime_fixtures",
"tests.helpers.unit.agent_runtime_fixtures",
"tests.helpers.unit.model_provider_fixtures",
)
# set global mock for Redis client
redis_mock = MagicMock()
redis_mock.get = MagicMock(return_value=None)
redis_mock.setex = MagicMock()
redis_mock.setnx = MagicMock()
redis_mock.delete = MagicMock()
redis_mock.lock = MagicMock()
redis_mock.exists = MagicMock(return_value=False)
redis_mock.set = MagicMock()
redis_mock.expire = MagicMock()
redis_mock.hgetall = MagicMock(return_value={})
redis_mock.hdel = MagicMock()
redis_mock.incr = MagicMock(return_value=1)
# pytest 9.2's monkeypatch undo deletes a mock child that isn't in the mock's __dict__. Keep every Redis
# method there, so a test that patches one leaves it in place for the tests after it.
for _name in {*dir(ext_redis.RedisClientWrapper), *dir(Redis)}:
if not _name.startswith("_"):
setattr(redis_mock, _name, getattr(redis_mock, _name)) # guard-ignore: no-new-getattr -- names come from dir()
if TYPE_CHECKING:
from extensions.application_services.app import AppServices
from services.tag_application_service import TagApplicationService
def _patch_redis_clients_on_loaded_modules() -> None:
"""Ensure any module-level redis_client references point to the shared redis_mock."""
import sys
for module in list(sys.modules.values()):
if module is None:
continue
for client_attribute in ("redis_client", "_pubsub_redis_client"):
if hasattr(module, client_attribute):
setattr(module, client_attribute, redis_mock)
@pytest.fixture
def redis_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[tuple[ext_redis.RedisClientWrapper, MagicMock]]:
"""Exercise the wrapper and Redis command builders with network dispatch replaced."""
apply_config_overrides(monkeypatch, REDIS_KEY_PREFIX="")
with Redis() as client, patch.object(client, "execute_command", return_value=None) as commands:
wrapper = ext_redis.RedisClientWrapper()
wrapper.initialize(client)
yield wrapper, commands
@pytest.fixture
def tenant_queue_commands(
redis_transport: tuple[ext_redis.RedisClientWrapper, MagicMock], monkeypatch: pytest.MonkeyPatch
) -> MagicMock:
"""Run tenant queue serialization and command building without Redis I/O."""
from core.rag.pipeline import queue
redis, commands = redis_transport
monkeypatch.setattr(queue, "redis_client", redis)
return commands
@pytest.fixture
def app() -> Flask:
return CACHED_APP
@pytest.fixture(autouse=True)
def _provide_app_context(app: Flask) -> Iterator[None]:
with app.app_context():
yield
@pytest.fixture(autouse=True)
def _patch_redis_clients() -> Iterator[None]:
"""Patch and rebind loaded Redis clients to the shared mock for each unit test."""
with (
patch.object(ext_redis, "redis_client", redis_mock),
patch.object(ext_redis, "_pubsub_redis_client", redis_mock),
):
_patch_redis_clients_on_loaded_modules()
yield
@pytest.fixture(autouse=True)
def reset_redis_mock(_patch_redis_clients: None) -> None:
"""Reset the shared Redis mock after per-test client rebinding."""
# ``mock_calls`` records calls on every attached child and return value, so when it is empty the
# recursive reset has nothing to clear. Most tests never touch Redis; skip walking ~370 child mocks.
if redis_mock.mock_calls:
redis_mock.reset_mock()
# Restoring a monkeypatched method can leave it detached from the parent's reset traversal.
redis_mock.delete.reset_mock()
redis_mock.get.reset_mock()
redis_mock.setex.reset_mock()
redis_mock.setnx.reset_mock()
redis_mock.lock.reset_mock(return_value=True, side_effect=True)
redis_mock.exists.reset_mock()
redis_mock.set.reset_mock()
redis_mock.expire.reset_mock()
redis_mock.hgetall.reset_mock()
redis_mock.hdel.reset_mock()
redis_mock.incr.reset_mock()
redis_mock.get.return_value = None
redis_mock.setex.return_value = None
redis_mock.setnx.return_value = None
redis_mock.delete.return_value = None
redis_mock.exists.return_value = False
redis_mock.set.return_value = None
redis_mock.expire.return_value = None
redis_mock.hgetall.return_value = dict[bytes, bytes]()
redis_mock.hdel.return_value = None
redis_mock.incr.return_value = 1
@pytest.fixture(autouse=True)
def reset_secret_key(monkeypatch: pytest.MonkeyPatch) -> None:
"""Ensure SECRET_KEY-dependent logic sees an empty config value by default."""
apply_config_overrides(monkeypatch, SECRET_KEY="")
@pytest.fixture
def config_overrides(monkeypatch: pytest.MonkeyPatch) -> Callable[..., None]:
"""Temporarily override fields on the shared typed application config.
Application modules import the same config instance, so mutating known
field names keeps tests scoped without replacing that instance with an
unconstrained mock. ``monkeypatch`` restores every value after the test.
"""
def apply(**values: object) -> None:
apply_config_overrides(monkeypatch, **values)
return apply
# The callable that materializes each live ``_sqlite_engine``'s database file.
_sqlite_copiers: WeakKeyDictionary[Engine, Callable[[], None]] = WeakKeyDictionary()
def _copy_sqlite_template(template: Path, database_path: Path, lock: Lock) -> None:
"""Create ``database_path`` as a copy of ``template`` unless it already exists.
Idempotent and safe across threads sharing ``lock``: the copy happens at most once, so an
open connection never sees its file overwritten.
"""
with lock:
if not database_path.exists():
shutil.copyfile(template, database_path)
@pytest.fixture
def _sqlite_engine(_sqlite_database_template: Path) -> Iterator[Engine]:
"""Copy the schema into an isolated directory without pytest's numbered scan.
``tmp_path`` searches all preceding test directories for a free number. This
autouse dependency needs only a unique, disposable directory, including any
SQLite journal files, so keep it separate from test-owned ``tmp_path`` data.
Most tests never connect, so the schema is copied on the engine's first connection
rather than for every test. Code that opens the file without this engine must go
through ``sqlite_engine``, which copies it up front.
"""
with TemporaryDirectory(prefix="case-", dir=_sqlite_database_template.parent) as directory:
database_path = Path(directory) / "unit-tests.sqlite3"
engine = create_engine(URL.create("sqlite", database=str(database_path)))
lock = Lock()
def copy() -> None:
_copy_sqlite_template(_sqlite_database_template, database_path, lock)
event.listen(engine, "do_connect", lambda *_args: copy())
_sqlite_copiers[engine] = copy
try:
yield engine
finally:
engine.dispose()
@pytest.fixture(scope="session")
def _sqlite_database_template(tmp_path_factory: pytest.TempPathFactory) -> Path:
"""Create one empty full-schema SQLite database per pytest worker."""
database_path = tmp_path_factory.mktemp("sqlite-template") / "unit-tests.sqlite3"
engine = create_engine(URL.create("sqlite", database=str(database_path)))
try:
TypeBase.metadata.create_all(engine)
finally:
engine.dispose()
return database_path
@pytest.fixture(autouse=True)
def _sqlite_session_factory(
_sqlite_engine: Engine,
monkeypatch: pytest.MonkeyPatch,
) -> sessionmaker[Session]:
"""Bind all unit-test Sessions to the pristine full-schema SQLite database."""
factory = sessionmaker(bind=_sqlite_engine, expire_on_commit=False)
monkeypatch.setattr(session_factory_module, "_session_maker", factory)
return factory
@pytest.fixture
def _unbound_session_factory(
_sqlite_session_factory: sessionmaker[Session],
monkeypatch: pytest.MonkeyPatch,
) -> sessionmaker[Session]:
"""Create one unbound factory and install it as the global test factory."""
factory = sessionmaker()
monkeypatch.setattr(session_factory_module, "_session_maker", factory)
return factory
@pytest.fixture
def sqlite_engine(_sqlite_engine: Engine) -> Engine:
"""Expose the pristine full-schema SQLite engine to tests.
The database file exists on return, so tests may hand ``sqlite_engine.url`` to code
that opens it with an engine of its own.
"""
_sqlite_copiers[_sqlite_engine]()
return _sqlite_engine
@pytest.fixture
def sqlite_session_factory(_sqlite_session_factory: sessionmaker[Session]) -> sessionmaker[Session]:
"""Expose the shared SQLite session factory to tests."""
return _sqlite_session_factory
@pytest.fixture
def sqlite_session(_sqlite_session_factory: sessionmaker[Session]) -> Iterator[Session]:
"""Yield a session over the pristine full-schema SQLite database.
Legacy indirect model parameters remain accepted by pytest but are ignored.
Remove those decorators as their test files receive individual review.
"""
with _sqlite_session_factory() as session:
yield session
@pytest.fixture
def unbound_session_factory(_unbound_session_factory: sessionmaker[Session]) -> sessionmaker[Session]:
"""Expose an unbound factory for paths that must not require persistence."""
return _unbound_session_factory
@pytest.fixture
def unbound_session(_unbound_session_factory: sessionmaker[Session]) -> Iterator[Session]:
"""Yield an unbound Session for paths that must not require persistence.
Bind-requiring database access fails, while bind-free Session operations can
still succeed.
"""
with _unbound_session_factory() as session:
yield session
def persist_service_api_tenant_owner(session: Session, tenant: Tenant, owner: Account) -> TenantAccountJoin:
"""Persist the owner identity resolved by service-API app authentication.
The legacy name is retained temporarily for consumers on independent
conversion branches, but this helper no longer fabricates an execute result.
"""
membership = TenantAccountJoin(
tenant_id=tenant.id,
account_id=owner.id,
role=TenantAccountRole.OWNER,
)
owner._current_tenant = tenant
session.add_all([tenant, owner, membership])
session.commit()
return membership
def persist_service_api_dataset_owner(
session: Session,
tenant: Tenant,
tenant_account_join: TenantAccountJoin,
) -> None:
"""Persist the tenant-owner mapping resolved by dataset-token authentication."""
session.add_all([tenant, tenant_account_join])
session.commit()
@pytest.fixture
def account_domain(
sqlite_session_factory: sessionmaker[Session],
monkeypatch: pytest.MonkeyPatch,
config_overrides: Callable[..., None],
) -> AccountDomain:
from blinker import Signal
from enums import DeploymentEdition
from services.workspace import gateways
from tests.helpers.unit.account_domain import build_account_domain
config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, RBAC_ENABLED=False)
monkeypatch.setattr(gateways, "generate_key_pair", lambda _workspace_id: "public-key")
monkeypatch.setattr(gateways, "tenant_was_created", Signal())
return build_account_domain(sqlite_session_factory)
@pytest.fixture
def account_application_services(
sqlite_session_factory: sessionmaker[Session], account_domain: AccountDomain
) -> ApplicationServices:
from dataclasses import replace
from enums import DeploymentEdition
from extensions.ext_application_services import build_application_services
from extensions.ext_redis import RedisClientWrapper
services = build_application_services(
database_client=sqlite_session_factory,
deployment_edition=DeploymentEdition.COMMUNITY,
initialization_password="",
redis=create_autospec(RedisClientWrapper, instance=True),
)
return replace(
services,
accounts=replace(services.accounts, lifecycle=account_domain.accounts),
workspaces=replace(
services.workspaces,
members=account_domain.members,
provisioning=account_domain.provisioning,
invitations=account_domain.invitations,
),
)
@pytest.fixture
def rbac_domain(sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch) -> RBACDomain:
return build_rbac_domain(sqlite_session_factory, monkeypatch)
@pytest.fixture
def app_services(sqlite_session_factory: sessionmaker[Session], rbac_domain: RBACDomain) -> AppServices:
from unittest.mock import Mock
from extensions.application_services.app import build_app_services
from extensions.ext_application_services import _build_oauth_server_service
from extensions.ext_redis import redis_client
from repositories.app.console_repository import ConsoleAppRepository
from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository
from services.app_tracing_config_service import AppTracingConfigService
from services.recommended_app_package_service import RecommendedAppPackageService
return build_app_services(
database_client=sqlite_session_factory,
repository=ConsoleAppRepository(session_factory=sqlite_session_factory),
rbac_members=rbac_domain.rbac.members,
tracing_configs=AppTracingConfigService(
configs=SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory),
# Unit tests never reach a tracing vendor; an unexpected request fails the test.
provider=deployment_provider_configs(RecordingEgress([], grpc_target=None)),
),
oauth=_build_oauth_server_service(database_client=sqlite_session_factory, redis=redis_client),
recommended_packages=RecommendedAppPackageService(sources=Mock(), exporter=Mock()),
)
@pytest.fixture
def application_tags(sqlite_session_factory: sessionmaker[Session]) -> TagApplicationService:
from repositories.tag_repository import TagRepository
from services.tag_application_service import TagApplicationService
return TagApplicationService(tags=TagRepository(sqlite_session_factory))