Austin/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2024-11-16 11:42:29 +08:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

fixed test_model_manager.py

Browse Source 
Signed-off-by: root <root@localhost.localdomain>
This commit is contained in:
root 2024-11-11 18:23:10 +08:00
parent 6327d72bc8
commit ec0c4a25fc
2 changed files with 18 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
api/extensions/ext_celery.py
View File

@ -43,6 +43,10 @@ def init_app(app: Flask) -> Celery:
result_backend=dify_config.CELERY_RESULT_BACKEND, result_backend=dify_config.CELERY_RESULT_BACKEND,
broker_transport_options=broker_transport_options, broker_transport_options=broker_transport_options,
broker_connection_retry_on_startup=True, broker_connection_retry_on_startup=True,
worker_log_format=dify_config.LOG_FORMAT,
worker_task_log_format=dify_config.LOG_FORMAT,
worker_logfile=dify_config.LOG_FILE,
worker_hijack_root_logger=False,
) )
if dify_config.BROKER_USE_SSL: if dify_config.BROKER_USE_SSL:

33
api/tests/unit_tests/core/test_model_manager.py
View File

@ -1,9 +1,9 @@
from unittest.mock import MagicMock from unittest.mock import MagicMock, patch
import pytest import pytest
import redis import redis
from extensions.ext_redis import RedisClientWrapper
from extensions.ext_redis import redis_client
from core.entities.provider_entities import ModelLoadBalancingConfiguration from core.entities.provider_entities import ModelLoadBalancingConfiguration
from core.model_manager import LBModelManager from core.model_manager import LBModelManager
from core.model_runtime.entities.model_entities import ModelType from core.model_runtime.entities.model_entities import ModelType
@ -39,12 +39,10 @@ def lb_model_manager():
return lb_model_manager return lb_model_manager
class FakeRedisClient:
def __init__(self, host='localhost', port=6379, db=0):
self.client = redis.Redis(host=host, port=port, db=db)
def test_lb_model_manager_fetch_next(mocker, lb_model_manager): def test_lb_model_manager_fetch_next(mocker, lb_model_manager):
# initialize redis client
redis_client.initialize(redis.Redis())
assert len(lb_model_manager._load_balancing_configs) == 3 assert len(lb_model_manager._load_balancing_configs) == 3
config1 = lb_model_manager._load_balancing_configs[0] config1 = lb_model_manager._load_balancing_configs[0]
@ -62,16 +60,13 @@ def test_lb_model_manager_fetch_next(mocker, lb_model_manager):
start_index += 1 start_index += 1
return start_index return start_index
fake_redis_client = FakeRedisClient() with (
redis_client = RedisClientWrapper() patch.object(redis_client, "incr", side_effect=incr),
redis_client.initialize(fake_redis_client) patch.object(redis_client, "set", return_value=None),
patch.object(redis_client, "expire", return_value=None),
mocker.patch("redis.Redis.incr", side_effect=incr) ):
mocker.patch("redis.Redis.set", return_value=None) config = lb_model_manager.fetch_next()
mocker.patch("redis.Redis.expire", return_value=None) assert config == config2
config = lb_model_manager.fetch_next() config = lb_model_manager.fetch_next()
assert config == config2 assert config == config3
config = lb_model_manager.fetch_next()
assert config == config3