From 38fe432443b955fcd73bdbf84da2b601c4c20522 Mon Sep 17 00:00:00 2001
From: Xander Steenbrugge <xandersteenbruggex@gmail.com>
Date: Mon, 6 Jan 2025 16:50:05 +0100
Subject: [PATCH] add specific test option

---
 comfyui.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/comfyui.py b/comfyui.py
index 425fb55..84b43dc 100644
--- a/comfyui.py
+++ b/comfyui.py
@@ -63,13 +63,20 @@
 test_workflows = os.getenv("WORKFLOWS")
 root_workflows_folder = "../private_workflows" if os.getenv("PRIVATE") else "../workflows"
 test_all = True if os.getenv("TEST_ALL") else False
+specific_test = os.getenv("SPECIFIC_TEST") if os.getenv("SPECIFIC_TEST") else ""
 skip_tests = os.getenv("SKIP_TESTS")
 
+if test_all and specific_test:
+    print(f"WARNING: can't have both TEST_ALL and SPECIFIC_TEST at the same time...")
+    print(f"Running TEST_ALL instead")
+    specific_test = ""
+
 print("========================================")
 print(f"db: {db}")
 print(f"workspace: {workspace_name}")
 print(f"test_workflows: {test_workflows}")
 print(f"test_all: {test_all}")
+print(f"specific_test: {specific_test}")
 print(f"skip_tests: {skip_tests}")
 print("========================================")
 
@@ -224,6 +231,7 @@ def download_files(force_redownload=False):
     modal.Image.debian_slim(python_version="3.11")
     .env({"COMFYUI_PATH": "/root", "COMFYUI_MODEL_PATH": "/root/models"}) 
     .env({"TEST_ALL": os.getenv("TEST_ALL")})
+    .env({"SPECIFIC_TEST": os.getenv("SPECIFIC_TEST")})
     .apt_install("git", "git-lfs", "libgl1-mesa-glx", "libglib2.0-0", "libmagic1", "ffmpeg", "libegl1")
     .pip_install_from_pyproject(str(root_dir / "pyproject.toml"))
     .pip_install("diffusers==0.31.0")
@@ -319,10 +327,11 @@ def downloads(self):
             
     @modal.build()
     def test_workflows(self):
-        print(" ==== TESTING WORKFLOWS ====")
         if os.getenv("SKIP_TESTS"):
             print("Skipping tests")
             return
+            
+        print(" ==== TESTING WORKFLOWS ====")
         
         t1 = time.time()
         self._start()
@@ -348,6 +357,8 @@ def test_workflows(self):
             test_all = os.getenv("TEST_ALL", False)
             if test_all:
                 tests = glob.glob(f"/root/workspace/workflows/{workflow}/test*.json")
+            elif specific_test:
+                tests = [f"/root/workspace/workflows/{workflow}/{specific_test}"]
             else:
                 tests = [f"/root/workspace/workflows/{workflow}/test.json"]
             print(f"====> Running tests for {workflow}: ", tests)