diff --git a/src/codemodder/codemods/__init__.py b/src/codemodder/codemods/__init__.py index 07bb27c5..f5db9d32 100644 --- a/src/codemodder/codemods/__init__.py +++ b/src/codemodder/codemods/__init__.py @@ -53,10 +53,12 @@ def match_codemods(codemod_include: list, codemod_exclude: list) -> dict: name: codemod for codemod in DEFAULT_CODEMODS if (name := codemod.name()) not in codemod_exclude + and (_ := codemod.id()) not in codemod_exclude } return { name: codemod for codemod in DEFAULT_CODEMODS if (name := codemod.name()) in codemod_include + or (_ := codemod.id()) in codemod_include } diff --git a/tests/test_main.py b/tests/test_main.py index 7b5a2975..cea54f83 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -110,6 +110,20 @@ def test_no_codemods_to_run(self, mock_semgrep_run): assert exit_code == 0 mock_semgrep_run.assert_not_called() + @pytest.mark.parametrize("codemod", ["secure-random", "pixee:python/secure-random"]) + @mock.patch("codemodder.__main__.compile_results") + def test_run_codemod_name_or_id(self, mock_compile_results, codemod): + args = [ + "tests/samples/", + "--output", + "here.txt", + f"--codemod-include={codemod}", + ] + + exit_code = run(parse_args(args), args) + assert exit_code == 0 + mock_compile_results.assert_called() + @mock.patch("codemodder.semgrep.results_by_path_and_rule_id") @mock.patch("codemodder.__main__.semgrep_run", side_effect=semgrep_run) @mock.patch("codemodder.semgrep.subprocess.run")