From e031120dfe8cf0ba79bd20eb130d143594b4c19a Mon Sep 17 00:00:00 2001 From: Yoland Y <4950057+yoland68@users.noreply.github.com> Date: Thu, 19 Dec 2024 15:12:33 -0800 Subject: [PATCH] refactor: update security checks in publish command and add tests 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. --- comfy_cli/command/custom_nodes/command.py | 35 +++----- tests/comfy_cli/command/nodes/test_publish.py | 84 +++++++++++++++++++ 2 files changed, 95 insertions(+), 24 deletions(-) create mode 100644 tests/comfy_cli/command/nodes/test_publish.py diff --git a/comfy_cli/command/custom_nodes/command.py b/comfy_cli/command/custom_nodes/command.py index 05a15a26..b1264768 100644 --- a/comfy_cli/command/custom_nodes/command.py +++ b/comfy_cli/command/custom_nodes/command.py @@ -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 @@ -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: diff --git a/tests/comfy_cli/command/nodes/test_publish.py b/tests/comfy_cli/command/nodes/test_publish.py new file mode 100644 index 00000000..7425835c --- /dev/null +++ b/tests/comfy_cli/command/nodes/test_publish.py @@ -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