tests/fix: include all models required for tests as fixtures

Haidra-Org · Oct 3, 2024 · cc298b3 · cc298b3
1 parent c753355
commit cc298b3
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -122,13 +122,17 @@ def isolated_comfy_horde_instance(init_horde) -> Comfy_Horde:
 _sdxl_refined_model_name = "AlbedoBase XL (SDXL)"
 _stable_cascade_base_model_name = "Stable Cascade 1.0"
 _flux1_schnell_fp8_base_model_name = "Flux.1-Schnell fp8 (Compact)"
+_am_pony_xl_model_name = "AMPonyXL"
+_rev_animated_model_name = "Rev Animated"
 
 _all_model_names = [
     _testing_model_name,
     _sdxl_1_0_model_name,
     _sdxl_refined_model_name,
     _stable_cascade_base_model_name,
     _flux1_schnell_fp8_base_model_name,
+    _am_pony_xl_model_name,
+    _rev_animated_model_name,
 ]
 
 # !!!!
@@ -166,6 +170,18 @@ def flux1_schnell_fp8_base_model_name(shared_model_manager: type[SharedModelMana
     return _flux1_schnell_fp8_base_model_name
 
 
+@pytest.fixture(scope="session")
+def am_pony_xl_model_name(shared_model_manager: type[SharedModelManager]) -> str:
+    """The default AMPonyXL model name used for testing."""
+    return _am_pony_xl_model_name
+
+
+@pytest.fixture(scope="session")
+def rev_animated_model_name(shared_model_manager: type[SharedModelManager]) -> str:
+    """The default Rev Animated model name used for testing."""
+    return _rev_animated_model_name
+
+
 # !!!!
 # If you're adding a model name, follow the pattern and **add it to `_all_model_names`**
 # !!!!

diff --git a/tests/test_horde_lora.py b/tests/test_horde_lora.py
@@ -390,6 +390,7 @@ def test_text_to_image_lora_character_hires_fix(
         self,
         shared_model_manager: type[SharedModelManager],
         hordelib_instance: HordeLib,
+        rev_animated_model_name: str,
     ):
         assert shared_model_manager.manager.lora
 
@@ -435,7 +436,7 @@ def test_text_to_image_lora_character_hires_fix(
             ],
             "ddim_steps": 30,
             "n_iter": 1,
-            "model": "Rev Animated",
+            "model": rev_animated_model_name,
         }
 
         pil_image = hordelib_instance.basic_inference_single_image(data).image
@@ -454,6 +455,7 @@ def test_text_to_image_lora_character_hires_fix_sdxl(
         self,
         shared_model_manager: type[SharedModelManager],
         hordelib_instance: HordeLib,
+        am_pony_xl_model_name: str,
     ):
         assert shared_model_manager.manager.lora
 
@@ -488,7 +490,7 @@ def test_text_to_image_lora_character_hires_fix_sdxl(
             ],
             "ddim_steps": 12,
             "n_iter": 1,
-            "model": "AMPonyXL",
+            "model": am_pony_xl_model_name,
         }
 
         pil_image = hordelib_instance.basic_inference_single_image(data).image
@@ -1010,6 +1012,7 @@ def test_login_gated_lora(
         self,
         shared_model_manager: type[SharedModelManager],
         hordelib_instance: HordeLib,
+        sdxl_1_0_base_model_name: str,
     ):
         assert shared_model_manager.manager.lora
 
@@ -1044,7 +1047,7 @@ def test_login_gated_lora(
             "loras": [{"name": download_gated_lora_version_id, "model": 1.0, "clip": 1.0, "is_version": True}],
             "ddim_steps": 25,
             "n_iter": 1,
-            "model": "SDXL 1.0",
+            "model": sdxl_1_0_base_model_name,
         }
         ret = hordelib_instance.basic_inference_single_image(data)
         assert isinstance(ret, ResultingImageReturn)