Skip to content

Commit

Permalink
refactor: update security checks in publish command and add tests
Browse files Browse the repository at this point in the history
Replaced the previous security check implementation with a subprocess call to 'ruff' for better error handling and output. Added tests to verify behavior on security violations, missing 'ruff' installation, and successful publishing with a token. This enhances the reliability and maintainability of the publish process.
  • Loading branch information
yoland68 committed Dec 20, 2024
1 parent cc38cc8 commit e031120
Show file tree
Hide file tree
Showing 2 changed files with 95 additions and 24 deletions.
35 changes: 11 additions & 24 deletions comfy_cli/command/custom_nodes/command.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,6 @@

import typer
from rich import print
from ruff import ConfigurationError
from ruff.__main__ import find_ruff_python_files
from ruff.linter import lint_file
from typing_extensions import Annotated, List

from comfy_cli import logging, tracking, ui, utils
Expand Down Expand Up @@ -712,31 +709,21 @@ def publish(
# Run security checks first
typer.echo("Running security checks...")
try:
files = find_ruff_python_files(["."])
settings = {
"select": ["S102", "S307"],
"ignore": [],
"line-length": 88,
}

has_violations = False
for python_file in files:
violations = lint_file(python_file, settings=settings)
if violations:
has_violations = True
for violation in violations:
print(f"[red]{violation}[/red]")

if has_violations:
print("[red]Security issues found. Please fix them before publishing.[/red]")
# Run ruff check with security rules
cmd = ["ruff", "check", ".", "--select", "S102,S307"]
result = subprocess.run(cmd, capture_output=True, text=True)

if result.returncode != 0:
print("[red]Security issues found:[/red]")
print(result.stdout)
raise typer.Exit(code=1)

except ConfigurationError as e:
print(f"[red]Configuration error during security check: {e}[/red]")
# raise typer.Exit(code=1)
except FileNotFoundError:
print("[red]Ruff is not installed. Please install it with 'pip install ruff'[/red]")
raise typer.Exit(code=1)
except Exception as e:
print(f"[red]Error running security check: {e}[/red]")
# raise typer.Exit(code=1)
raise typer.Exit(code=1)

# Prompt for API Key
if not token:
Expand Down
84 changes: 84 additions & 0 deletions tests/comfy_cli/command/nodes/test_publish.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
from unittest.mock import MagicMock, patch

from typer.testing import CliRunner

from comfy_cli.command.custom_nodes.command import app

runner = CliRunner()


def test_publish_fails_on_security_violations():
# Mock subprocess.run to simulate security violations
mock_result = MagicMock()
mock_result.returncode = 1
mock_result.stdout = "S102 Use of exec() detected"

with patch("subprocess.run", return_value=mock_result):
result = runner.invoke(app, ["publish"])

assert result.exit_code == 1
assert "Security issues found" in result.stdout


def test_publish_continues_on_no_security_violations():
# Mock subprocess.run to simulate no violations
mock_result = MagicMock()
mock_result.returncode = 0
mock_result.stdout = ""

with (
patch("subprocess.run", return_value=mock_result),
patch("comfy_cli.command.custom_nodes.command.extract_node_configuration") as mock_extract,
patch("typer.prompt") as mock_prompt,
patch("comfy_cli.command.custom_nodes.command.registry_api.publish_node_version") as mock_publish,
patch("comfy_cli.command.custom_nodes.command.zip_files") as mock_zip,
patch("comfy_cli.command.custom_nodes.command.upload_file_to_signed_url") as mock_upload,
):
# Setup the mocks
mock_extract.return_value = {"name": "test-node"}
mock_prompt.return_value = "test-token"
mock_publish.return_value = MagicMock(signedUrl="https://test.url")

# Run the publish command
_result = runner.invoke(app, ["publish"])

# Verify the publish flow continued
assert mock_extract.called
assert mock_publish.called
assert mock_zip.called
assert mock_upload.called


def test_publish_handles_missing_ruff():
with patch("subprocess.run", side_effect=FileNotFoundError()):
result = runner.invoke(app, ["publish"])

assert result.exit_code == 1
assert "Ruff is not installed" in result.stdout


def test_publish_with_token_option():
# Mock subprocess.run to simulate no violations
mock_result = MagicMock()
mock_result.returncode = 0
mock_result.stdout = ""

with (
patch("subprocess.run", return_value=mock_result),
patch("comfy_cli.command.custom_nodes.command.extract_node_configuration") as mock_extract,
patch("comfy_cli.command.custom_nodes.command.registry_api.publish_node_version") as mock_publish,
patch("comfy_cli.command.custom_nodes.command.zip_files") as mock_zip,
patch("comfy_cli.command.custom_nodes.command.upload_file_to_signed_url") as mock_upload,
):
# Setup the mocks
mock_extract.return_value = {"name": "test-node"}
mock_publish.return_value = MagicMock(signedUrl="https://test.url")

# Run the publish command with token
_result = runner.invoke(app, ["publish", "--token", "test-token"])

# Verify the publish flow worked with provided token
assert mock_extract.called
assert mock_publish.called
assert mock_zip.called
assert mock_upload.called

0 comments on commit e031120

Please sign in to comment.