ada e2e test (#138)

Co-authored-by: James Kwon <[email protected]>

.github/workflows/build-and-test.yml

@@ -13,7 +13,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     env:
       PYTHONIOENCODING: "utf8"
-    
+
     strategy:
       fail-fast: false
       matrix:
@@ -31,11 +31,13 @@ jobs:
 
       - name: Install Dependencies
         run: |
+          python -m pip install --upgrade pip
+          pip install pytest
           pip install -e .
 
-      - name: Run Core Functionality
+      - name: Test e2e
+        env:
+          PYTHONPATH: ${{ github.workspace }}
+          TEST_E2E: true
         run: |
-          comfy --skip-prompt --no-enable-telemetry env
-          comfy --skip-prompt install --cpu
-          comfy launch --background -- --cpu
-          
+          pytest tests/e2e

.github/workflows/run-on-gpu.yml

@@ -5,25 +5,24 @@ on:
     branches:
       - main
     paths:
-      - 'comfy_cli/**'
-      - '!comfy_cli/test_**'
+      - "comfy_cli/**"
+      - "!comfy_cli/test_**"
   pull_request:
     branches:
       - main
     paths:
-      - 'comfy_cli/**'
-      - '!comfy_cli/test_**'
+      - "comfy_cli/**"
+      - "!comfy_cli/test_**"
 
 jobs:
   test-cli-gpu:
     name: "Run Tests on GPU Runners"
-    runs-on: 
+    runs-on:
       group: gpu-runners
-      labels: ${{ matrix.os }}-x64-gpu # 
+      labels: ${{ matrix.os }}-x64-gpu #
     strategy:
       fail-fast: false
       matrix:
-        # windows test is flaky.
         os: [linux]
 
     steps:
@@ -39,16 +38,30 @@ jobs:
         with:
           python-version: 3.12
 
+      - name: Check disk space
+        run: |
+          df -h
+
       - name: Install Dependencies
         run: |
+          python -m pip install --upgrade pip
+          pip install pytest
           pip install -e .
 
-      - name: Check disk space
+      - name: Test e2e
+        id: test-e2e
+        env:
+          PYTHONPATH: ${{ github.workspace }}
+          TEST_E2E: true
+          TEST_E2E_COMFY_INSTALL_FLAGS: --nvidia --cuda-version 11.8
+          TEST_E2E_COMFY_LAUNCH_FLAGS_EXTRA: ""
         run: |
-          df -h
+          pytest tests/e2e
 
-      - name: Run Core Functionality
+      - name: Retry test e2e but without gpu
+        if: ${{ failure() && steps.test-e2e.conclusion == 'failure' }}
+        env:
+          PYTHONPATH: ${{ github.workspace }}
+          TEST_E2E: true
         run: |
-          comfy --skip-prompt --no-enable-telemetry env
-          comfy --skip-prompt install ${{ matrix.comfy-flags }} --nvidia --cuda-version 11.8
-          comfy launch --background
+          pytest tests/e2e

.pre-commit-config.yaml

@@ -14,4 +14,4 @@ repos:
   - repo: https://github.com/PyCQA/isort
     rev: 5.13.2
     hooks:
-      - id: isort
+      - id: isort

comfy_cli/command/custom_nodes/cm_cli_util.py

@@ -53,6 +53,7 @@ def execute_cm_cli(args, channel=None, mode=None) -> str | None:
         result = subprocess.run(
             cmd, env=new_env, check=True, capture_output=True, text=True
         )
+        print(result.stdout)
         return result.stdout
     except subprocess.CalledProcessError as e:
         if e.returncode == 1:

comfy_cli/command/models/models.py

@@ -145,6 +145,13 @@ def download(
             show_default=True,
         ),
     ] = None,
+    filename: Annotated[
+        Optional[str],
+        typer.Option(
+            help="The filename to save the model.",
+            show_default=True,
+        ),
+    ] = None,
     set_civitai_api_token: Annotated[
         Optional[str],
         typer.Option(
@@ -226,9 +233,15 @@ def download(
     else:
         print("Model source is unknown")
 
-    local_filename = ui.prompt_input(
-        "Enter filename to save model as", default=local_filename
-    )
+    if filename is None:
+        if local_filename is None:
+            local_filename = ui.prompt_input("Enter filename to save model as")
+        else:
+            local_filename = ui.prompt_input(
+                "Enter filename to save model as", default=local_filename
+            )
+    else:
+        local_filename = filename
 
     if relative_path is None:
         relative_path = DEFAULT_COMFY_MODEL_PATH
@@ -264,6 +277,11 @@ def remove(
         help="List of model filenames to delete, separated by spaces",
         show_default=False,
     ),
+    confirm: bool = typer.Option(
+        False,
+        help="Confirm for deletion and skip the prompt",
+        show_default=False,
+    ),
 ):
     """Remove one or more downloaded models, either by specifying them directly or through an interactive selection."""
     model_dir = get_workspace() / relative_path
@@ -304,8 +322,11 @@ def remove(
         to_delete = [model_dir / selection for selection in selections]
 
     # Confirm deletion
-    if to_delete and ui.prompt_confirm_action(
-        "Are you sure you want to delete the selected files?", False
+    if to_delete and (
+        confirm
+        or ui.prompt_confirm_action(
+            "Are you sure you want to delete the selected files?", False
+        )
     ):
         for model_path in to_delete:
             model_path.unlink()

comfy_cli/registry/api.py

@@ -13,8 +13,11 @@ def __init__(self):
         self.base_url = self.determine_base_url()
 
     def determine_base_url(self):
-        if os.getenv("ENVIRONMENT") == "dev":
+        env = os.getenv("ENVIRONMENT")
+        if env == "dev":
             return "http://localhost:8080"
+        elif env == "staging":
+            return "https://staging.comfyregistry.org"
         else:
             return "https://api.comfy.org"
 

pyproject.toml

@@ -56,4 +56,4 @@ include = ["comfy_cli*"]
 
 [tool.isort]
 profile = "black"
-line_length = 88
+line_length = 88

tests/e2e/test_e2e.py

@@ -0,0 +1,173 @@
+import os
+import subprocess
+from datetime import datetime
+from textwrap import dedent
+
+import pytest
+
+
+def e2e_test(func):
+    return pytest.mark.skipif(
+        os.getenv("TEST_E2E", "false") != "true",
+        reason="Test e2e is not explicitly enabled",
+    )(func)
+
+
+def exec(cmd: str, **kwargs) -> subprocess.CompletedProcess[str]:
+    cmd = dedent(cmd).strip()
+    print(f"cmd: {cmd}")
+
+    proc = subprocess.run(
+        args=cmd,
+        capture_output=True,
+        text=True,
+        shell=True,
+        encoding="utf-8",
+        check=False,
+        **kwargs,
+    )
+    print(proc.stdout, proc.stderr)
+    return proc
+
+
+@pytest.fixture(scope="module")
+def workspace():
+    ws = os.path.join(os.getcwd(), f"comfy-{datetime.now().timestamp()}")
+    proc = exec(
+        f"""
+            comfy --skip-prompt --workspace {ws} install {os.getenv("TEST_E2E_COMFY_INSTALL_FLAGS", "--cpu")}
+            comfy --skip-prompt set-default {ws}
+            comfy --skip-prompt --no-enable-telemetry env
+        """
+    )
+    assert 0 == proc.returncode
+
+    proc = exec(
+        f"""
+        comfy --workspace {ws} launch --background -- {os.getenv("TEST_E2E_COMFY_LAUNCH_FLAGS_EXTRA", "--cpu")}
+        """
+    )
+    assert 0 == proc.returncode
+
+    yield ws
+
+    proc = exec(
+        f"""
+        comfy --workspace {ws} stop
+        """
+    )
+    assert 0 == proc.returncode
+
+
+@pytest.fixture()
+def comfy_cli(workspace):
+    exec("comfy --skip-prompt --no-enable-telemetry env")
+    return f"comfy --workspace {workspace}"
+
+
+@e2e_test
+def test_model(comfy_cli):
+    url = "https://huggingface.co/guoyww/animatediff/resolve/cd71ae134a27ec6008b968d6419952b0c0494cf2/mm_sd_v14.ckpt?download=true"
+    path = os.path.join("models", "animatediff_models")
+    proc = exec(
+        f"""
+            {comfy_cli} model download --url {url} --relative-path {path} --filename animatediff_models
+        """
+    )
+    assert 0 == proc.returncode
+
+    proc = exec(
+        f"""
+            {comfy_cli} model list --relative-path {path}
+        """
+    )
+    assert 0 == proc.returncode
+    assert "animatediff_models" in proc.stdout
+
+    proc = exec(
+        f"""
+            {comfy_cli} model remove --relative-path {path} --model-names animatediff_models --confirm
+        """
+    )
+    assert 0 == proc.returncode
+
+
+@e2e_test
+def test_node(comfy_cli, workspace):
+    node = "ComfyUI-AnimateDiff-Evolved"
+    proc = exec(
+        f"""
+            {comfy_cli} node install {node}
+        """
+    )
+    assert 0 == proc.returncode
+
+    proc = exec(
+        f"""
+            {comfy_cli} node reinstall {node}
+        """
+    )
+    assert 0 == proc.returncode
+
+    proc = exec(
+        f"""
+            {comfy_cli} node show all
+        """
+    )
+    assert 0 == proc.returncode
+    assert node in proc.stdout
+
+    proc = exec(
+        f"""
+            {comfy_cli} node update {node}
+        """
+    )
+    assert 0 == proc.returncode
+
+    proc = exec(
+        f"""
+            {comfy_cli} node disable {node}
+        """
+    )
+    assert 0 == proc.returncode
+
+    proc = exec(
+        f"""
+            {comfy_cli} node enable {node}
+        """
+    )
+    assert 0 == proc.returncode
+
+    pubID = "comfytest123"
+    pubToken = "6075cf7b-47e7-4c58-a3de-38f59a9bcc22"
+    proc = exec(
+        f"""
+            sed 's/PublisherId = ".*"/PublisherId = "{pubID}"/g' pyproject.toml
+            {comfy_cli} node publish --token {pubToken}
+        """,
+        env={"ENVIRONMENT": "stage"},
+        cwd=os.path.join(workspace, "custom_nodes", node),
+    )
+
+
+@e2e_test
+def test_run(comfy_cli):
+    url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt"
+    path = os.path.join("models", "checkpoints")
+    name = "v1-5-pruned-emaonly.ckpt"
+    proc = exec(
+        f"""
+            {comfy_cli} model download --url {url} --relative-path {path} --filename {name}
+        """
+    )
+    assert 0 == proc.returncode
+
+    workflow = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "workflow.json"
+    )
+    proc = exec(
+        f"""
+        {comfy_cli} run --workflow {workflow} --wait --timeout 180
+        """
+    )
+    assert 0 == proc.returncode