From 6b03388f76597f745cf30e24f6cbd1306ab9049f Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Tue, 10 Jun 2025 18:05:20 +0530 Subject: [PATCH] fix test --- tests/test_inference_manager.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_inference_manager.py b/tests/test_inference_manager.py index 62bd509..28bf43f 100644 --- a/tests/test_inference_manager.py +++ b/tests/test_inference_manager.py @@ -16,17 +16,17 @@ def before_all() -> None: def test_get_inference_pool() -> None: - model_names = [ 'yolo_nsfw' ] - model_source_set = content_analyser.get_model_options().get('sources') + model_names = [ 'yolo_11m', 'marqo', 'freepik' ] + _, model_source_set = content_analyser.collect_model_downloads() with patch('facefusion.inference_manager.detect_app_context', return_value = 'cli'): get_inference_pool('facefusion.content_analyser', model_names, model_source_set) - assert isinstance(INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.yolo_nsfw.0.cpu').get('content_analyser'), InferenceSession) + assert isinstance(INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.yolo_11m.marqo.freepik.0.cpu').get('yolo_11m'), InferenceSession) with patch('facefusion.inference_manager.detect_app_context', return_value = 'ui'): get_inference_pool('facefusion.content_analyser', model_names, model_source_set) - assert isinstance(INFERENCE_POOL_SET.get('ui').get('facefusion.content_analyser.yolo_nsfw.0.cpu').get('content_analyser'), InferenceSession) + assert isinstance(INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.yolo_11m.marqo.freepik.0.cpu').get('yolo_11m'), InferenceSession) - assert INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.yolo_nsfw.0.cpu').get('content_analyser') == INFERENCE_POOL_SET.get('ui').get('facefusion.content_analyser.yolo_nsfw.0.cpu').get('content_analyser') + assert INFERENCE_POOL_SET.get('cli').get('facefusion.content_analyser.yolo_11m.marqo.freepik.0.cpu').get('yolo_11m') == INFERENCE_POOL_SET.get('ui').get('facefusion.content_analyser.yolo_11m.marqo.freepik.0.cpu').get('yolo_11m')