jprdonnelly/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-05-25 19:00:43 -04:00
Code Issues Packages Projects Releases Wiki Activity

fix: use Generator type annotation with @contextmanager decorators (#36297)

Browse Source 
Co-authored-by: Asuka Minato <i@asukaminato.eu.org>
This commit is contained in:
EvanYao
2026-05-18 15:12:49 +08:00
committed by GitHub
parent 969760364d
commit d1238180ed
3 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
api/tests/unit_tests/core/app/test_llm_quota.py
View File

@@ -1,4 +1,4 @@
from collections.abc import Iterator
from collections.abc import Generator
from contextlib import contextmanager
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
@@ -26,7 +26,7 @@ from models.provider import Provider, ProviderType
@contextmanager
def _patched_credit_pool_session_factory(engine: Engine) -> Iterator[None]:
def _patched_credit_pool_session_factory(engine: Engine) -> Generator[None, None, None]:
session_maker = sessionmaker(bind=engine, expire_on_commit=False)
with patch("services.credit_pool_service.session_factory.get_session_maker", return_value=session_maker):
yield

4
api/tests/unit_tests/events/test_update_provider_when_message_created.py
View File

@@ -1,4 +1,4 @@
from collections.abc import Iterator
from collections.abc import Generator
from contextlib import contextmanager
from types import SimpleNamespace
from unittest.mock import patch
@@ -16,7 +16,7 @@ from models.provider import ProviderType
@contextmanager
def _patched_credit_pool_session_factory(engine: Engine) -> Iterator[None]:
def _patched_credit_pool_session_factory(engine: Engine) -> Generator[None, None, None]:
session_maker = sessionmaker(bind=engine, expire_on_commit=False)
with patch("services.credit_pool_service.session_factory.get_session_maker", return_value=session_maker):
yield

4
api/tests/unit_tests/services/test_credit_pool_service.py
View File

@@ -1,4 +1,4 @@
from collections.abc import Iterator
from collections.abc import Generator
from contextlib import contextmanager
from unittest.mock import patch
from uuid import uuid4
@@ -34,7 +34,7 @@ def _create_engine_with_pool(*, quota_limit: int, quota_used: int) -> tuple[Engi
@contextmanager
def _patched_session_factory(engine: Engine) -> Iterator[None]:
def _patched_session_factory(engine: Engine) -> Generator[None, None, None]:
session_maker = sessionmaker(bind=engine, expire_on_commit=False)
with patch("services.credit_pool_service.session_factory.get_session_maker", return_value=session_maker):
yield