From ec0c4a25fc64dc7ddd6880b8ab523c17c361b571 Mon Sep 17 00:00:00 2001
From: root <root@localhost.localdomain>
Date: Mon, 11 Nov 2024 18:23:10 +0800
Subject: [PATCH] fixed test_model_manager.py

Signed-off-by: root <root@localhost.localdomain>
---
 api/extensions/ext_celery.py                  |  4 +++
 .../unit_tests/core/test_model_manager.py     | 33 ++++++++-----------
 2 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py
index 504899c276..52e14cc49b 100644
--- a/api/extensions/ext_celery.py
+++ b/api/extensions/ext_celery.py
@@ -43,6 +43,10 @@ def init_app(app: Flask) -> Celery:
         result_backend=dify_config.CELERY_RESULT_BACKEND,
         broker_transport_options=broker_transport_options,
         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:
diff --git a/api/tests/unit_tests/core/test_model_manager.py b/api/tests/unit_tests/core/test_model_manager.py
index 96176bb6c4..e37533e520 100644
--- a/api/tests/unit_tests/core/test_model_manager.py
+++ b/api/tests/unit_tests/core/test_model_manager.py
@@ -1,9 +1,9 @@
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 import redis
-from extensions.ext_redis import RedisClientWrapper
 
+from extensions.ext_redis import redis_client
 from core.entities.provider_entities import ModelLoadBalancingConfiguration
 from core.model_manager import LBModelManager
 from core.model_runtime.entities.model_entities import ModelType
@@ -39,12 +39,10 @@ def 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):
+    # initialize redis client
+    redis_client.initialize(redis.Redis())
+
     assert len(lb_model_manager._load_balancing_configs) == 3
 
     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
         return start_index
 
-    fake_redis_client = FakeRedisClient()
-    redis_client = RedisClientWrapper()
-    redis_client.initialize(fake_redis_client)
-    
-    mocker.patch("redis.Redis.incr", side_effect=incr)
-    mocker.patch("redis.Redis.set", return_value=None)
-    mocker.patch("redis.Redis.expire", return_value=None)
+    with (
+        patch.object(redis_client, "incr", side_effect=incr),
+        patch.object(redis_client, "set", return_value=None),
+        patch.object(redis_client, "expire", return_value=None),
+    ):
+        config = lb_model_manager.fetch_next()
+        assert config == config2
 
-    config = lb_model_manager.fetch_next()
-    assert config == config2
-
-    config = lb_model_manager.fetch_next()
-    assert config == config3
+        config = lb_model_manager.fetch_next()
+        assert config == config3