diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 7a3e705f16..c164f4c442 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -295,18 +295,17 @@ def test_heuristic_import_with_type(mm2_installer: ModelInstallServiceBase, mode "name": f"{model_params['name']}_2", "type": ModelType(model_params["type"]), } - try: - assert "repo_id" in model_params - install_job1 = mm2_installer.heuristic_import(source=model_params["repo_id"], config=config1) + assert "repo_id" in model_params + install_job1 = mm2_installer.heuristic_import(source=model_params["repo_id"], config=config1) + mm2_installer.wait_for_installs(timeout=10) + if model_params["type"] != "embedding": + assert install_job1.errored + assert install_job1.error_type == 'InvalidModelConfigException' + return + assert install_job1.complete + assert install_job1.config_out if model_params["type"] == "embedding" else not install_job1.config_out - while not install_job1.in_terminal_state: - sleep(0.01) - assert install_job1.config_out if model_params["type"] == "embedding" else not install_job1.config_out - - install_job2 = mm2_installer.heuristic_import(source=model_params["repo_id"], config=config2) - - while not install_job2.in_terminal_state: - sleep(0.01) - assert install_job2.config_out if model_params["type"] == "embedding" else not install_job2.config_out - except InvalidModelConfigException: - assert model_params["type"] != "embedding" + install_job2 = mm2_installer.heuristic_import(source=model_params["repo_id"], config=config2) + mm2_installer.wait_for_installs(timeout=10) + assert install_job2.complete + assert install_job2.config_out if model_params["type"] == "embedding" else not install_job2.config_out