import unittest from unittest.mock import patch, mock_open, MagicMock, call import json import sys import os import argparse # Mock data for runtimes-matrix.json mock_runtimes_matrix = [ { "name": "dev", "package": "kitchensink-runtime", "path": "substrate/frame", "header": "substrate/HEADER-APACHE2", "template": "substrate/.maintain/frame-weight-template.hbs", "bench_features": "runtime-benchmarks", "bench_flags": "--flag1 --flag2" }, { "name": "zagros", "package": "zagros-runtime", "path": "pezkuwi/runtime/zagros", "header": "pezkuwi/file_header.txt", "template": "pezkuwi/xcm/pallet-xcm-benchmarks/template.hbs", "bench_features": "runtime-benchmarks", "bench_flags": "--flag3 --flag4" }, { "name": "pezkuwichain", "package": "pezkuwichain-runtime", "path": "pezkuwi/runtime/pezkuwichain", "header": "pezkuwi/file_header.txt", "template": "pezkuwi/xcm/pallet-xcm-benchmarks/template.hbs", "bench_features": "runtime-benchmarks", "bench_flags": "" }, { "name": "asset-hub-zagros", "package": "asset-hub-zagros-runtime", "path": "cumulus/teyrchains/runtimes/assets/asset-hub-zagros", "header": "cumulus/file_header.txt", "template": "cumulus/templates/xcm-bench-template.hbs", "bench_features": "runtime-benchmarks", "bench_flags": "--flag7 --flag8" } ] def get_mock_bench_output(runtime, pallets, output_path, header, bench_flags, template = None): return f"pezframe-omni-bencher v1 benchmark pezpallet --extrinsic=* " \ f"--runtime=target/production/wbuild/{runtime}-runtime/{runtime.replace('-', '_')}_runtime.wasm " \ f"--pezpallet={pallets} --header={header} " \ f"--output={output_path} " \ f"--wasm-execution=compiled " \ f"--steps=50 --repeat=20 --heap-pages=4096 " \ f"{f'--template={template} ' if template else ''}" \ f"--no-storage-info --no-min-squares --no-median-slopes " \ f"{bench_flags}" class TestCmd(unittest.TestCase): def setUp(self): self.patcher1 = patch('builtins.open', new_callable=mock_open, read_data=json.dumps(mock_runtimes_matrix)) self.patcher2 = patch('json.load', return_value=mock_runtimes_matrix) self.patcher3 = patch('argparse.ArgumentParser.parse_known_args') self.patcher4 = patch('os.system', return_value=0) self.patcher5 = patch('os.popen') self.patcher6 = patch('importlib.util.spec_from_file_location', return_value=MagicMock()) self.patcher7 = patch('importlib.util.module_from_spec', return_value=MagicMock()) self.patcher8 = patch('cmd.generate_prdoc.main', return_value=0) self.mock_open = self.patcher1.start() self.mock_json_load = self.patcher2.start() self.mock_parse_args = self.patcher3.start() self.mock_system = self.patcher4.start() self.mock_popen = self.patcher5.start() self.mock_spec_from_file_location = self.patcher6.start() self.mock_module_from_spec = self.patcher7.start() self.mock_generate_prdoc_main = self.patcher8.start() # Ensure that cmd.py uses the mock_runtimes_matrix import cmd cmd.runtimesMatrix = mock_runtimes_matrix def tearDown(self): self.patcher1.stop() self.patcher2.stop() self.patcher3.stop() self.patcher4.stop() self.patcher5.stop() self.patcher6.stop() self.patcher7.stop() self.patcher8.stop() def test_bench_command_normal_execution_all_runtimes(self): self.mock_parse_args.return_value = (argparse.Namespace( command='bench-omni', runtime=list(map(lambda x: x['name'], mock_runtimes_matrix)), pezpallet=['pallet_balances'], fail_fast=True, quiet=False, clean=False, image=None ), []) self.mock_popen.return_value.read.side_effect = [ "pallet_balances\npallet_staking\npallet_something\n", # Output for dev runtime "pallet_balances\npallet_staking\npallet_something\n", # Output for zagros runtime "pallet_staking\npallet_something\n", # Output for pezkuwichain runtime - no pallet here "pallet_balances\npallet_staking\npallet_something\n", # Output for asset-hub-zagros runtime "./substrate/frame/balances/Cargo.toml\n", # Mock manifest path for dev -> pallet_balances ] with patch('sys.exit') as mock_exit: import cmd cmd.main() mock_exit.assert_not_called() expected_calls = [ # Build calls call("forklift cargo build -q -p kitchensink-runtime --profile production --features=runtime-benchmarks"), call("forklift cargo build -q -p zagros-runtime --profile production --features=runtime-benchmarks"), call("forklift cargo build -q -p pezkuwichain-runtime --profile production --features=runtime-benchmarks"), call("forklift cargo build -q -p asset-hub-zagros-runtime --profile production --features=runtime-benchmarks"), call(get_mock_bench_output( runtime='kitchensink', pallets='pallet_balances', output_path='./substrate/frame/balances/src/weights.rs', header=os.path.abspath('substrate/HEADER-APACHE2'), bench_flags='--flag1 --flag2', template="substrate/.maintain/frame-weight-template.hbs" )), call(get_mock_bench_output( runtime='zagros', pallets='pallet_balances', output_path='./pezkuwi/runtime/zagros/src/weights', header=os.path.abspath('pezkuwi/file_header.txt'), bench_flags='--flag3 --flag4' )), # skips pezkuwichain benchmark call(get_mock_bench_output( runtime='asset-hub-zagros', pallets='pallet_balances', output_path='./cumulus/teyrchains/runtimes/assets/asset-hub-zagros/src/weights', header=os.path.abspath('cumulus/file_header.txt'), bench_flags='--flag7 --flag8' )), ] self.mock_system.assert_has_calls(expected_calls, any_order=True) def test_bench_command_normal_execution(self): self.mock_parse_args.return_value = (argparse.Namespace( command='bench-omni', runtime=['zagros'], pezpallet=['pallet_balances', 'pallet_staking'], fail_fast=True, quiet=False, clean=False, image=None ), []) header_path = os.path.abspath('pezkuwi/file_header.txt') self.mock_popen.return_value.read.side_effect = [ "pallet_balances\npallet_staking\npallet_something\n", # Output for zagros runtime ] with patch('sys.exit') as mock_exit: import cmd cmd.main() mock_exit.assert_not_called() expected_calls = [ # Build calls call("forklift cargo build -q -p zagros-runtime --profile production --features=runtime-benchmarks"), # Zagros runtime calls call(get_mock_bench_output( runtime='zagros', pallets='pallet_balances', output_path='./pezkuwi/runtime/zagros/src/weights', header=header_path, bench_flags='--flag3 --flag4' )), call(get_mock_bench_output( runtime='zagros', pallets='pallet_staking', output_path='./pezkuwi/runtime/zagros/src/weights', header=header_path, bench_flags='--flag3 --flag4' )), ] self.mock_system.assert_has_calls(expected_calls, any_order=True) def test_bench_command_normal_execution_xcm(self): self.mock_parse_args.return_value = (argparse.Namespace( command='bench-omni', runtime=['zagros'], pezpallet=['pallet_xcm_benchmarks::generic'], fail_fast=True, quiet=False, clean=False, image=None ), []) header_path = os.path.abspath('pezkuwi/file_header.txt') self.mock_popen.return_value.read.side_effect = [ "pallet_balances\npallet_staking\npallet_something\npallet_xcm_benchmarks::generic\n", # Output for zagros runtime ] with patch('sys.exit') as mock_exit: import cmd cmd.main() mock_exit.assert_not_called() expected_calls = [ # Build calls call("forklift cargo build -q -p zagros-runtime --profile production --features=runtime-benchmarks"), # Zagros runtime calls call(get_mock_bench_output( runtime='zagros', pallets='pallet_xcm_benchmarks::generic', output_path='./pezkuwi/runtime/zagros/src/weights/xcm', header=header_path, bench_flags='--flag3 --flag4', template="pezkuwi/xcm/pallet-xcm-benchmarks/template.hbs" )), ] self.mock_system.assert_has_calls(expected_calls, any_order=True) def test_bench_command_two_runtimes_two_pallets(self): self.mock_parse_args.return_value = (argparse.Namespace( command='bench-omni', runtime=['zagros', 'pezkuwichain'], pezpallet=['pallet_balances', 'pallet_staking'], fail_fast=True, quiet=False, clean=False, image=None ), []) self.mock_popen.return_value.read.side_effect = [ "pallet_staking\npallet_balances\n", # Output for zagros runtime "pallet_staking\npallet_balances\n", # Output for pezkuwichain runtime ] with patch('sys.exit') as mock_exit: import cmd cmd.main() mock_exit.assert_not_called() header_path = os.path.abspath('pezkuwi/file_header.txt') expected_calls = [ # Build calls call("forklift cargo build -q -p zagros-runtime --profile production --features=runtime-benchmarks"), call("forklift cargo build -q -p pezkuwichain-runtime --profile production --features=runtime-benchmarks"), # Zagros runtime calls call(get_mock_bench_output( runtime='zagros', pallets='pallet_staking', output_path='./pezkuwi/runtime/zagros/src/weights', header=header_path, bench_flags='--flag3 --flag4' )), call(get_mock_bench_output( runtime='zagros', pallets='pallet_balances', output_path='./pezkuwi/runtime/zagros/src/weights', header=header_path, bench_flags='--flag3 --flag4' )), # Pezkuwichain runtime calls call(get_mock_bench_output( runtime='pezkuwichain', pallets='pallet_staking', output_path='./pezkuwi/runtime/pezkuwichain/src/weights', header=header_path, bench_flags='' )), call(get_mock_bench_output( runtime='pezkuwichain', pallets='pallet_balances', output_path='./pezkuwi/runtime/pezkuwichain/src/weights', header=header_path, bench_flags='' )), ] self.mock_system.assert_has_calls(expected_calls, any_order=True) def test_bench_command_one_dev_runtime(self): self.mock_parse_args.return_value = (argparse.Namespace( command='bench-omni', runtime=['dev'], pezpallet=['pallet_balances'], fail_fast=True, quiet=False, clean=False, image=None ), []) manifest_dir = "substrate/frame/kitchensink" self.mock_popen.return_value.read.side_effect = [ "pallet_balances\npallet_something", # Output for dev runtime manifest_dir + "/Cargo.toml" # Output for manifest path in dev runtime ] header_path = os.path.abspath('substrate/HEADER-APACHE2') with patch('sys.exit') as mock_exit: import cmd cmd.main() mock_exit.assert_not_called() expected_calls = [ # Build calls call("forklift cargo build -q -p kitchensink-runtime --profile production --features=runtime-benchmarks"), # Westend runtime calls call(get_mock_bench_output( runtime='kitchensink', pallets='pallet_balances', output_path=manifest_dir + "/src/weights.rs", header=header_path, bench_flags='--flag1 --flag2', template="substrate/.maintain/frame-weight-template.hbs" )), ] self.mock_system.assert_has_calls(expected_calls, any_order=True) def test_bench_command_one_cumulus_runtime(self): self.mock_parse_args.return_value = (argparse.Namespace( command='bench-omni', runtime=['asset-hub-zagros'], pezpallet=['pallet_assets'], fail_fast=True, quiet=False, clean=False, image=None ), []) self.mock_popen.return_value.read.side_effect = [ "pallet_assets\n", # Output for asset-hub-zagros runtime ] header_path = os.path.abspath('cumulus/file_header.txt') with patch('sys.exit') as mock_exit: import cmd cmd.main() mock_exit.assert_not_called() expected_calls = [ # Build calls call("forklift cargo build -q -p asset-hub-zagros-runtime --profile production --features=runtime-benchmarks"), # Asset-hub-zagros runtime calls call(get_mock_bench_output( runtime='asset-hub-zagros', pallets='pallet_assets', output_path='./cumulus/teyrchains/runtimes/assets/asset-hub-zagros/src/weights', header=header_path, bench_flags='--flag7 --flag8' )), ] self.mock_system.assert_has_calls(expected_calls, any_order=True) def test_bench_command_one_cumulus_runtime_xcm(self): self.mock_parse_args.return_value = (argparse.Namespace( command='bench-omni', runtime=['asset-hub-zagros'], pezpallet=['pallet_xcm_benchmarks::generic', 'pallet_assets'], fail_fast=True, quiet=False, clean=False, image=None ), []) self.mock_popen.return_value.read.side_effect = [ "pallet_assets\npallet_xcm_benchmarks::generic\n", # Output for asset-hub-zagros runtime ] header_path = os.path.abspath('cumulus/file_header.txt') with patch('sys.exit') as mock_exit: import cmd cmd.main() mock_exit.assert_not_called() expected_calls = [ # Build calls call("forklift cargo build -q -p asset-hub-zagros-runtime --profile production --features=runtime-benchmarks"), # Asset-hub-zagros runtime calls call(get_mock_bench_output( runtime='asset-hub-zagros', pallets='pallet_xcm_benchmarks::generic', output_path='./cumulus/teyrchains/runtimes/assets/asset-hub-zagros/src/weights/xcm', header=header_path, bench_flags='--flag7 --flag8', template="cumulus/templates/xcm-bench-template.hbs" )), call(get_mock_bench_output( runtime='asset-hub-zagros', pallets='pallet_assets', output_path='./cumulus/teyrchains/runtimes/assets/asset-hub-zagros/src/weights', header=header_path, bench_flags='--flag7 --flag8' )), ] self.mock_system.assert_has_calls(expected_calls, any_order=True) @patch('argparse.ArgumentParser.parse_known_args', return_value=(argparse.Namespace(command='fmt'), [])) @patch('os.system', return_value=0) def test_fmt_command(self, mock_system, mock_parse_args): with patch('sys.exit') as mock_exit: import cmd cmd.main() mock_exit.assert_not_called() mock_system.assert_any_call('cargo +nightly fmt') mock_system.assert_any_call('taplo format --config .config/taplo.toml') @patch('argparse.ArgumentParser.parse_known_args', return_value=(argparse.Namespace(command='update-ui'), [])) @patch('os.system', return_value=0) def test_update_ui_command(self, mock_system, mock_parse_args): with patch('sys.exit') as mock_exit: import cmd cmd.main() mock_exit.assert_not_called() mock_system.assert_called_with('sh ./scripts/update-ui-tests.sh') @patch('argparse.ArgumentParser.parse_known_args', return_value=(argparse.Namespace(command='prdoc'), [])) @patch('os.system', return_value=0) def test_prdoc_command(self, mock_system, mock_parse_args): with patch('sys.exit') as mock_exit: import cmd cmd.main() mock_exit.assert_not_called() self.mock_generate_prdoc_main.assert_called_with(mock_parse_args.return_value[0]) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'true', 'IS_PR_AUTHOR': 'false', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_valid_labels(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command with valid labels""" mock_get_labels.return_value = ['T1-FRAME', 'R0-no-crate-publish-required', 'D2-substantial'] mock_check_pr_status.return_value = True # PR is open mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['T1-FRAME', 'R0-no-crate-publish-required'] ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_not_called() # Check that JSON output was printed json_call = None for call in mock_print.call_args_list: if 'LABELS_JSON:' in str(call): json_call = call break self.assertIsNotNone(json_call) self.assertIn('T1-FRAME', str(json_call)) self.assertIn('R0-no-crate-publish-required', str(json_call)) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'true', 'IS_PR_AUTHOR': 'false', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_auto_correction(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command with auto-correctable typos""" mock_get_labels.return_value = ['T1-FRAME', 'R0-no-crate-publish-required', 'D2-substantial'] mock_check_pr_status.return_value = True # PR is open mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['T1-FRAM', 'R0-no-crate-publish'] # Typos that should be auto-corrected ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_not_called() # Check for auto-correction messages correction_messages = [str(call) for call in mock_print.call_args_list if 'Auto-corrected' in str(call)] self.assertTrue(len(correction_messages) > 0) # Check that JSON output contains corrected labels json_call = None for call in mock_print.call_args_list: if 'LABELS_JSON:' in str(call): json_call = call break self.assertIsNotNone(json_call) self.assertIn('T1-FRAME', str(json_call)) self.assertIn('R0-no-crate-publish-required', str(json_call)) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'true', 'IS_PR_AUTHOR': 'false', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_prefix_correction(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command with prefix matching""" mock_get_labels.return_value = ['T1-FRAME', 'T2-pallets', 'R0-no-crate-publish-required'] mock_check_pr_status.return_value = True # PR is open mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['T1-something'] # Should match T1-FRAME as the only T1- label ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_not_called() # Check that JSON output contains corrected label json_call = None for call in mock_print.call_args_list: if 'LABELS_JSON:' in str(call): json_call = call break self.assertIsNotNone(json_call) self.assertIn('T1-FRAME', str(json_call)) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'true', 'IS_PR_AUTHOR': 'false', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_invalid_labels(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command with invalid labels that cannot be corrected""" mock_get_labels.return_value = ['T1-FRAME', 'R0-no-crate-publish-required', 'D2-substantial'] mock_check_pr_status.return_value = True # PR is open mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['INVALID-LABEL', 'ANOTHER-BAD-LABEL'] ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_called_with(1) # Should exit with error code # Check for error JSON output error_json_call = None for call in mock_print.call_args_list: if 'ERROR_JSON:' in str(call): error_json_call = call break self.assertIsNotNone(error_json_call) self.assertIn('validation_failed', str(error_json_call)) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'true', 'IS_PR_AUTHOR': 'false', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_mixed_valid_invalid(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command with mix of valid and invalid labels""" mock_get_labels.return_value = ['T1-FRAME', 'R0-no-crate-publish-required', 'D2-substantial'] mock_check_pr_status.return_value = True # PR is open mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['T1-FRAME', 'INVALID-LABEL', 'D2-substantial'] ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_called_with(1) # Should exit with error code due to invalid label # Check for error JSON output error_json_call = None for call in mock_print.call_args_list: if 'ERROR_JSON:' in str(call): error_json_call = call break self.assertIsNotNone(error_json_call) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'true', 'IS_PR_AUTHOR': 'false', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_fetch_failure(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command when label fetching fails""" mock_get_labels.side_effect = RuntimeError("Failed to fetch labels from repository. Please check your connection and try again.") mock_check_pr_status.return_value = True # PR is open mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['T1-FRAME'] ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_called_with(1) # Should exit with error code # Check for error JSON output error_json_call = None for call in mock_print.call_args_list: if 'ERROR_JSON:' in str(call): error_json_call = call break self.assertIsNotNone(error_json_call) self.assertIn('Failed to fetch labels from repository', str(error_json_call)) def test_auto_correct_labels_function(self): """Test the auto_correct_labels function directly""" import cmd valid_labels = ['T1-FRAME', 'R0-no-crate-publish-required', 'D2-substantial', 'I2-bug'] # Test high similarity auto-correction corrections, suggestions = cmd.auto_correct_labels(['T1-FRAM'], valid_labels) self.assertEqual(len(corrections), 1) self.assertEqual(corrections[0][0], 'T1-FRAM') self.assertEqual(corrections[0][1], 'T1-FRAME') # Test low similarity suggestions corrections, suggestions = cmd.auto_correct_labels(['TOTALLY-WRONG'], valid_labels) self.assertEqual(len(corrections), 0) self.assertEqual(len(suggestions), 1) def test_find_closest_labels_function(self): """Test the find_closest_labels function directly""" import cmd valid_labels = ['T1-FRAME', 'T2-pallets', 'R0-no-crate-publish-required'] # Test finding close matches matches = cmd.find_closest_labels('T1-FRAM', valid_labels) self.assertIn('T1-FRAME', matches) # Test no close matches matches = cmd.find_closest_labels('COMPLETELY-DIFFERENT', valid_labels, cutoff=0.8) self.assertEqual(len(matches), 0) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'true', 'IS_PR_AUTHOR': 'false', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_merged_pr(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command on merged PR should fail""" mock_get_labels.return_value = ['T1-FRAME', 'R0-no-crate-publish-required'] mock_check_pr_status.return_value = False # PR is merged/closed mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['T1-FRAME'] ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_called_with(1) # Check for error JSON output error_json_call = None for call in mock_print.call_args_list: if 'ERROR_JSON:' in str(call): error_json_call = call break self.assertIsNotNone(error_json_call) self.assertIn('Cannot modify labels on merged PRs', str(error_json_call)) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'true', 'IS_PR_AUTHOR': 'false', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_open_pr(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command on open PR should succeed""" mock_get_labels.return_value = ['T1-FRAME', 'R0-no-crate-publish-required'] mock_check_pr_status.return_value = True # PR is open mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['T1-FRAME'] ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_not_called() # Check that JSON output was printed json_call = None for call in mock_print.call_args_list: if 'LABELS_JSON:' in str(call): json_call = call break self.assertIsNotNone(json_call) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'false', 'IS_PR_AUTHOR': 'false', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_unauthorized_user(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command by unauthorized user should fail""" mock_get_labels.return_value = ['T1-FRAME', 'R0-no-crate-publish-required'] mock_check_pr_status.return_value = True # PR is open mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['T1-FRAME'] ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_called_with(1) # Check for error JSON output error_json_call = None for call in mock_print.call_args_list: if 'ERROR_JSON:' in str(call): error_json_call = call break self.assertIsNotNone(error_json_call) self.assertIn('Only the PR author or organization members can modify labels', str(error_json_call)) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'false', 'IS_PR_AUTHOR': 'true', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_pr_author(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command by PR author should succeed""" mock_get_labels.return_value = ['T1-FRAME', 'R0-no-crate-publish-required'] mock_check_pr_status.return_value = True # PR is open mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['T1-FRAME'] ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_not_called() # Check that JSON output was printed json_call = None for call in mock_print.call_args_list: if 'LABELS_JSON:' in str(call): json_call = call break self.assertIsNotNone(json_call) @patch.dict('os.environ', {'PR_NUM': '123', 'IS_ORG_MEMBER': 'true', 'IS_PR_AUTHOR': 'false', 'GITHUB_TOKEN': 'fake_token'}) @patch('cmd.get_allowed_labels') @patch('cmd.check_pr_status') @patch('argparse.ArgumentParser.parse_known_args') def test_label_command_org_member(self, mock_parse_args, mock_check_pr_status, mock_get_labels): """Test label command by org member should succeed""" mock_get_labels.return_value = ['T1-FRAME', 'R0-no-crate-publish-required'] mock_check_pr_status.return_value = True # PR is open mock_parse_args.return_value = (argparse.Namespace( command='label', labels=['T1-FRAME'] ), []) with patch('sys.exit') as mock_exit, patch('builtins.print') as mock_print: import cmd cmd.main() mock_exit.assert_not_called() # Check that JSON output was printed json_call = None for call in mock_print.call_args_list: if 'LABELS_JSON:' in str(call): json_call = call break self.assertIsNotNone(json_call) if __name__ == '__main__': unittest.main()