diff --git a/services/hwfit/image_models.py b/services/hwfit/image_models.py index 088a00b..f47b602 100644 --- a/services/hwfit/image_models.py +++ b/services/hwfit/image_models.py @@ -288,7 +288,7 @@ def rank_image_models(system, search=None, sort="fit"): for model in IMAGE_MODEL_REGISTRY: # Filter by search - if search: + if isinstance(search, str) and search: s = search.lower() if s not in model["name"].lower() and s not in model["id"].lower() and s not in model.get("description", "").lower(): continue diff --git a/tests/test_image_models_nonstring_search.py b/tests/test_image_models_nonstring_search.py new file mode 100644 index 0000000..830e4ae --- /dev/null +++ b/tests/test_image_models_nonstring_search.py @@ -0,0 +1,15 @@ +from services.hwfit.image_models import rank_image_models, IMAGE_MODEL_REGISTRY + +SYS = {"gpu_vram_gb": 0, "has_gpu": False} + + +def test_rank_image_models_handles_non_string_search(): + # search is a CLI/API filter arg; a non-string made search.lower() raise + # AttributeError. A non-string search should behave as "no filter". + out = rank_image_models(SYS, search=123) + assert len(out) == len(IMAGE_MODEL_REGISTRY) + + +def test_rank_image_models_string_filter_still_applies(): + out = rank_image_models(SYS, search="zzzznotarealmodelzzz") + assert out == []