From 5375ed5fa9b0c2cd640f362775f052df4860b2fa Mon Sep 17 00:00:00 2001 From: Ignatious Johnson Date: Mon, 11 May 2026 19:22:42 +0000 Subject: [PATCH 1/2] Add upload_file and download_file support to MultiProcessPssh with operation registry This commit addresses API incompatibility issues between Pssh and MultiProcessPssh classes and implements a robust operation registry system to prevent future drift. Key Changes: - Add missing upload_file and download_file methods to MultiProcessPssh class - Implement operation registry (SUPPORTED_OPERATIONS) as single source of truth - Use method names directly as operation names to eliminate mapping errors - Add missing detailed parameter support to MultiProcessPssh.exec() API Parity Testing: - Add comprehensive test suite to catch missing methods automatically - Validate method signatures match between Pssh and MultiProcessPssh - Verify operation registry completeness and consistency - Test runtime execution of all sharder operations with proper parameters Operations Updated: - exec -> exec (added detailed parameter support) - upload_file -> upload_file (new operation, matches method name) - download_file -> download_file (new operation, matches method name) This ensures MultiProcessPssh maintains full API compatibility with Pssh and prevents the class of issues that led to missing upload_file/download_file methods. Co-authored-by: Cursor Signed-off-by: Ignatious Johnson --- cvs/lib/parallel/multiprocess_pssh.py | 60 +++- cvs/lib/parallel/pssh_sharder.py | 113 ++++++- .../unittests/test_multiprocess_pssh.py | 104 +++++- .../parallel/unittests/test_pssh_sharder.py | 298 +++++++++++++++++- 4 files changed, 548 insertions(+), 27 deletions(-) diff --git a/cvs/lib/parallel/multiprocess_pssh.py b/cvs/lib/parallel/multiprocess_pssh.py index 8e2e8c4b..7b47ffcd 100644 --- a/cvs/lib/parallel/multiprocess_pssh.py +++ b/cvs/lib/parallel/multiprocess_pssh.py @@ -171,10 +171,10 @@ def _print_merged_outputs(self, cmd_output, cmd=None, cmd_list=None, print_conso for line in cmd_output[host].splitlines(): self.log.info("%s", line) - def exec(self, cmd, timeout=None, print_console=True): + def exec(self, cmd, timeout=None, print_console=True, detailed=False): """Execute command with automatic sharding if needed.""" if not getattr(self, '_use_process_sharding', False): - return super().exec(cmd, timeout=timeout, print_console=print_console) + return super().exec(cmd, timeout=timeout, print_console=print_console, detailed=detailed) if self.env_prefix: full_cmd = f"{self.env_prefix} ; {cmd}" @@ -195,7 +195,7 @@ def exec(self, cmd, timeout=None, print_console=True): # Use sharder for sharded execution host_chunks = list(self.sharder.chunk_hosts(self.reachable_hosts)) payloads = self.sharder.create_payloads( - 'exec', host_chunks, self._shard_init_kwargs(), cmd=cmd, timeout=timeout + 'exec', host_chunks, self._shard_init_kwargs(), cmd=cmd, timeout=timeout, detailed=detailed ) shard_returns = self.sharder.execute_sharded(payloads) @@ -256,7 +256,7 @@ def exec_cmd_list(self, cmd_list, timeout=None, print_console=True): offset += k payloads.append( self.sharder.create_payloads( - 'cmd_list', [shard_hosts], self._shard_init_kwargs(), cmd_list=shard_commands, timeout=timeout + 'exec_cmd_list', [shard_hosts], self._shard_init_kwargs(), cmd_list=shard_commands, timeout=timeout )[0] ) @@ -283,7 +283,7 @@ def scp_file(self, local_file, remote_file, recurse=False): host_chunks = list(self.sharder.chunk_hosts(self.reachable_hosts)) payloads = self.sharder.create_payloads( - 'scp', + 'scp_file', host_chunks, self._shard_init_kwargs(), local_file=local_file, @@ -293,6 +293,54 @@ def scp_file(self, local_file, remote_file, recurse=False): shard_returns = self.sharder.execute_sharded(payloads) return self._merge_shard_returns(shard_returns, merge_unreachable=False) + def upload_file(self, local_file, remote_file, recurse=False): + """Upload file with automatic sharding if needed.""" + if not getattr(self, '_use_process_sharding', False): + return super().upload_file(local_file, remote_file, recurse=recurse) + + self.log.info('SFTP upload %s -> %s on %s', local_file, remote_file, self.reachable_hosts) + + host_chunks = list(self.sharder.chunk_hosts(self.reachable_hosts)) + payloads = self.sharder.create_payloads( + 'upload_file', + host_chunks, + self._shard_init_kwargs(), + local_file=local_file, + remote_file=remote_file, + recurse=recurse, + ) + shard_returns = self.sharder.execute_sharded(payloads) + return self._merge_shard_returns(shard_returns, merge_unreachable=False) + + def download_file(self, remote_file, local_file, recurse=False, suffix_separator='_'): + """Download file with automatic sharding if needed.""" + if not getattr(self, '_use_process_sharding', False): + return super().download_file(remote_file, local_file, recurse=recurse, suffix_separator=suffix_separator) + + self.log.info('SFTP download %s -> %s from %s', remote_file, local_file, self.reachable_hosts) + + host_chunks = list(self.sharder.chunk_hosts(self.reachable_hosts)) + payloads = self.sharder.create_payloads( + 'download_file', + host_chunks, + self._shard_init_kwargs(), + remote_file=remote_file, + local_file=local_file, + recurse=recurse, + suffix_separator=suffix_separator, + ) + shard_returns = self.sharder.execute_sharded(payloads) + + # For download_file, we need to merge the returned dictionaries from each shard + # The result is a dict mapping {host: actual_local_path} + merged_result = {} + for shard in shard_returns: + result = shard.get('result') or {} + if isinstance(result, dict): + merged_result.update(result) + + return merged_result + def reboot_connections(self): """Reboot connections with automatic sharding if needed.""" if not getattr(self, '_use_process_sharding', False): @@ -301,7 +349,7 @@ def reboot_connections(self): self.log.info('Rebooting Connections') host_chunks = list(self.sharder.chunk_hosts(self.reachable_hosts)) - payloads = self.sharder.create_payloads('reboot', host_chunks, self._shard_init_kwargs()) + payloads = self.sharder.create_payloads('reboot_connections', host_chunks, self._shard_init_kwargs()) shard_returns = self.sharder.execute_sharded(payloads) return self._merge_shard_returns(shard_returns, merge_unreachable=False) diff --git a/cvs/lib/parallel/pssh_sharder.py b/cvs/lib/parallel/pssh_sharder.py index b1f27994..f15c9b19 100644 --- a/cvs/lib/parallel/pssh_sharder.py +++ b/cvs/lib/parallel/pssh_sharder.py @@ -11,6 +11,42 @@ from cvs.lib.parallel.pssh import Pssh +# Operation registry - single source of truth for supported operations +# Operation names now match method names for consistency and simplicity +SUPPORTED_OPERATIONS = { + 'exec': { + 'required_params': ['cmd'], + 'optional_params': ['timeout', 'detailed'], + 'returns_result': True, + }, + 'exec_cmd_list': { + 'required_params': ['cmd_list'], + 'optional_params': ['timeout'], + 'returns_result': True, + }, + 'scp_file': { + 'required_params': ['local_file', 'remote_file'], + 'optional_params': ['recurse'], + 'returns_result': False, + }, + 'upload_file': { + 'required_params': ['local_file', 'remote_file'], + 'optional_params': ['recurse'], + 'returns_result': False, + }, + 'download_file': { + 'required_params': ['remote_file', 'local_file'], + 'optional_params': ['recurse', 'suffix_separator'], + 'returns_result': True, + }, + 'reboot_connections': { + 'required_params': [], + 'optional_params': [], + 'returns_result': False, + }, +} + + class PsshSharder: """Shards SSH operations across multiple processes for large host lists.""" @@ -90,23 +126,74 @@ def run_shard(payload): shard = Pssh(**init_kwargs) try: - # Direct operation calls - no registry needed! + # Use operation registry for parameter validation and documentation + op_config = SUPPORTED_OPERATIONS.get(operation) + if not op_config: + # Let the else clause in the operation handler catch unknown operations + op_config = {'required_params': [], 'optional_params': [], 'returns_result': False} + method_name = operation # Operation name IS the method name! + shard_method = getattr(shard, method_name) + + # Build arguments from payload based on operation configuration + args = {} + + # Add required parameters + for param in op_config['required_params']: + if param not in payload: + raise ValueError(f"Missing required parameter '{param}' for operation '{operation}'") + args[param] = payload[param] + + # Add optional parameters if present + for param in op_config['optional_params']: + if param in payload: + args[param] = payload[param] + + # Add common parameters for operations that support them + if method_name in ['exec', 'exec_cmd_list']: + args['print_console'] = False # Always False for sharded operations + + # Execute the operation - handle special cases for method signatures if operation == 'exec': - result = shard.exec(payload['cmd'], timeout=payload.get('timeout'), print_console=False) - elif operation == 'cmd_list': - result = shard.exec_cmd_list( - payload['cmd_list'], - timeout=payload.get('timeout'), - print_console=False, - ) - elif operation == 'scp': - shard.scp_file(payload['local_file'], payload['remote_file'], recurse=payload.get('recurse', False)) + # Maintain existing call signature for backward compatibility + kwargs = { + 'timeout': args.get('timeout'), + 'print_console': False, + } + if 'detailed' in args: + kwargs['detailed'] = args['detailed'] + result = shard_method(args['cmd'], **kwargs) + elif operation == 'exec_cmd_list': + # Maintain existing call signature for backward compatibility + result = shard_method(args['cmd_list'], timeout=args.get('timeout'), print_console=False) + elif operation in ['scp_file', 'upload_file']: + # File upload operations + shard_method(args['local_file'], args['remote_file'], recurse=args.get('recurse', False)) result = None - elif operation == 'reboot': - shard.reboot_connections() + elif operation == 'download_file': + # File download operation + result = shard_method( + args['remote_file'], + args['local_file'], + recurse=args.get('recurse', False), + suffix_separator=args.get('suffix_separator', '_'), + ) + elif operation == 'reboot_connections': + # Reboot operation + shard_method() result = None else: - raise ValueError(f'Unknown operation: {operation}') + # Enhanced error handling to distinguish user vs implementation errors + if operation in SUPPORTED_OPERATIONS: + # Case #2: Implementation Error - Operation in registry but no handler + raise RuntimeError( + f'Implementation bug: Operation \'{operation}\' is supported (in registry) ' + f'but handler is missing in run_shard().' + ) + else: + # Case #1: User Error - Operation not supported at all + raise ValueError( + f'Unknown operation: {operation}. Supported operations: {list(SUPPORTED_OPERATIONS.keys())}' + ) return { 'result': result, diff --git a/cvs/lib/parallel/unittests/test_multiprocess_pssh.py b/cvs/lib/parallel/unittests/test_multiprocess_pssh.py index eec7a8d2..6896a07f 100644 --- a/cvs/lib/parallel/unittests/test_multiprocess_pssh.py +++ b/cvs/lib/parallel/unittests/test_multiprocess_pssh.py @@ -109,7 +109,7 @@ def test_exec_no_sharding(self): pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test") result = pssh.exec("uptime") - mock_parent_exec.assert_called_once_with("uptime", timeout=None, print_console=True) + mock_parent_exec.assert_called_once_with("uptime", timeout=None, print_console=True, detailed=False) self.assertEqual(result, {"host1": "result"}) def test_exec_with_sharding(self): @@ -405,6 +405,104 @@ def test_merge_shard_returns(self): self.assertEqual(set(pssh.reachable_hosts), {"host1", "host2"}) self.assertEqual(pssh.unreachable_hosts, []) + def test_upload_file_with_sharding(self): + """Test upload_file with sharding uses sharder.""" + config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) + + with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): + with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: + mock_sharder = MagicMock() + mock_sharder_class.return_value = mock_sharder + + # Set up the mock sharder behavior + mock_sharder.chunk_hosts.return_value = [["host1"], ["host2"]] + mock_sharder.create_payloads.return_value = [ + {"operation": "upload", "init": {}, "local_file": "test.txt", "remote_file": "/tmp/test.txt"}, + ] + mock_sharder.execute_sharded.return_value = [ + { + "result": None, # upload_file returns None on success + "reachable_hosts": ["host1", "host2"], + "unreachable_hosts": [], + }, + ] + + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + + # Manually set attributes that _init_sharded would set + pssh.host_list = self.host_list + pssh.reachable_hosts = self.host_list + pssh.sharder = mock_sharder + pssh._use_process_sharding = True # Enable sharding for this test + # Add attributes needed by _shard_init_kwargs() + pssh.user = "test" + pssh.password = None + pssh.pkey = "id_rsa" + pssh.host_key_check = False + pssh.stop_on_errors = True + pssh.env_vars = None + + result = pssh.upload_file("test.txt", "/tmp/test.txt", recurse=True) + + # Verify sharder methods were called + mock_sharder.chunk_hosts.assert_called_once_with(self.host_list) + mock_sharder.execute_sharded.assert_called_once() + # upload_file should return the merged results (empty dict for success) + self.assertEqual(result, {}) + + def test_download_file_with_sharding(self): + """Test download_file with sharding uses sharder and merges results.""" + config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) -if __name__ == "__main__": - unittest.main() + with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): + with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: + mock_sharder = MagicMock() + mock_sharder_class.return_value = mock_sharder + + # Set up the mock sharder behavior + mock_sharder.chunk_hosts.return_value = [["host1"], ["host2"]] + mock_sharder.create_payloads.return_value = [ + { + "operation": "download", + "init": {}, + "remote_file": "/remote/test.txt", + "local_file": "/tmp/test.txt", + }, + ] + # Simulate download results from each shard + mock_sharder.execute_sharded.return_value = [ + { + "result": {"host1": "/tmp/test.txt_host1"}, # Download results from shard 1 + "reachable_hosts": ["host1"], + "unreachable_hosts": [], + }, + { + "result": {"host2": "/tmp/test.txt_host2"}, # Download results from shard 2 + "reachable_hosts": ["host2"], + "unreachable_hosts": [], + }, + ] + + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + + # Manually set attributes that _init_sharded would set + pssh.host_list = self.host_list + pssh.reachable_hosts = self.host_list + pssh.sharder = mock_sharder + pssh._use_process_sharding = True # Enable sharding for this test + # Add attributes needed by _shard_init_kwargs() + pssh.user = "test" + pssh.password = None + pssh.pkey = "id_rsa" + pssh.host_key_check = False + pssh.stop_on_errors = True + pssh.env_vars = None + + result = pssh.download_file("/remote/test.txt", "/tmp/test.txt", suffix_separator="_") + + # Verify sharder methods were called + mock_sharder.chunk_hosts.assert_called_once_with(self.host_list) + mock_sharder.execute_sharded.assert_called_once() + # download_file should return merged host->path dict + expected_result = {"host1": "/tmp/test.txt_host1", "host2": "/tmp/test.txt_host2"} + self.assertEqual(result, expected_result) diff --git a/cvs/lib/parallel/unittests/test_pssh_sharder.py b/cvs/lib/parallel/unittests/test_pssh_sharder.py index 8c8bbf89..13f37752 100644 --- a/cvs/lib/parallel/unittests/test_pssh_sharder.py +++ b/cvs/lib/parallel/unittests/test_pssh_sharder.py @@ -78,12 +78,12 @@ def test_create_payloads_cmd_list_mode(self): shard_init_kwargs = {'user': 'testuser'} payloads = self.sharder.create_payloads( - 'cmd_list', host_chunks, shard_init_kwargs, cmd_list=['echo 1', 'echo 2'] + 'exec_cmd_list', host_chunks, shard_init_kwargs, cmd_list=['echo 1', 'echo 2'] ) expected = [ { - 'operation': 'cmd_list', + 'operation': 'exec_cmd_list', 'init': {'user': 'testuser', 'host_list': ['host1', 'host2']}, 'cmd_list': ['echo 1', 'echo 2'], } @@ -243,7 +243,7 @@ def test_pssh_shard_worker_cmd_list_mode(self, mock_pssh_class): mock_shard.unreachable_hosts = [] payload = { - 'operation': 'cmd_list', + 'operation': 'exec_cmd_list', 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, 'cmd_list': ['echo 1', 'echo 2'], } @@ -267,7 +267,7 @@ def test_pssh_shard_worker_scp_mode(self, mock_pssh_class): mock_shard.unreachable_hosts = [] payload = { - 'operation': 'scp', + 'operation': 'scp_file', 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, 'local_file': 'test.txt', 'remote_file': '/tmp/test.txt', @@ -291,7 +291,7 @@ def test_pssh_shard_worker_reboot_mode(self, mock_pssh_class): mock_shard.unreachable_hosts = [] payload = { - 'operation': 'reboot', + 'operation': 'reboot_connections', 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, } @@ -319,6 +319,294 @@ def test_pssh_shard_worker_unknown_operation(self, mock_pssh_class): self.assertIn('Unknown operation: unknown_operation', str(cm.exception)) + @patch('cvs.lib.parallel.pssh_sharder.Pssh') + def test_pssh_shard_worker_registry_handler_mismatch(self, mock_pssh_class): + """Test shard worker with operation in registry but missing handler (implementation error).""" + mock_shard = MagicMock() + mock_pssh_class.return_value = mock_shard + + # Temporarily add an operation to registry without corresponding handler + from cvs.lib.parallel.pssh_sharder import SUPPORTED_OPERATIONS + + original_operations = SUPPORTED_OPERATIONS.copy() + + try: + # Add a fake operation to registry + SUPPORTED_OPERATIONS['fake_operation'] = { + 'required_params': [], + 'optional_params': [], + 'returns_result': False, + } + + payload = { + 'operation': 'fake_operation', + 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, + } + + with self.assertRaises(RuntimeError) as cm: + PsshSharder.run_shard(payload) + + # Verify it's identified as a implementation error + self.assertIn('Implementation bug', str(cm.exception)) + self.assertIn('fake_operation', str(cm.exception)) + self.assertIn('supported (in registry) but handler is missing', str(cm.exception)) + + finally: + # Restore original registry + SUPPORTED_OPERATIONS.clear() + SUPPORTED_OPERATIONS.update(original_operations) + + +class TestPsshSharderApiCompliance(unittest.TestCase): + """ + Test suite to ensure PsshSharder supports all operations and parameters + required by MultiProcessPssh, preventing API drift issues. + """ + + def test_sharder_supports_all_multiprocess_operations(self): + """ + Verify that pssh_sharder supports all operations used by MultiProcessPssh. + + Uses the operation registry as the single source of truth instead of + fragile regex parsing. + """ + import re + import inspect + from cvs.lib.parallel.pssh_sharder import SUPPORTED_OPERATIONS + + # Get MultiProcessPssh source code to extract sharder operations + from cvs.lib.parallel.multiprocess_pssh import MultiProcessPssh + + multiprocess_source = inspect.getsource(MultiProcessPssh) + + # Extract operation names from create_payloads calls + operation_pattern = r"create_payloads\s*\(\s*['\"](\w+)['\"]" + used_operations = set(re.findall(operation_pattern, multiprocess_source)) + + # Get supported operations from registry (single source of truth) + supported_operations = set(SUPPORTED_OPERATIONS.keys()) + + # Find missing operations + missing_operations = used_operations - supported_operations + + if missing_operations: + error_msg = ( + f"PsshSharder missing support for {len(missing_operations)} operation(s) " + f"used by MultiProcessPssh:\n" + f"Missing operations: {sorted(missing_operations)}\n" + f"Used by MultiProcessPssh: {sorted(used_operations)}\n" + f"Supported by PsshSharder: {sorted(supported_operations)}\n\n" + f"Add missing operations to SUPPORTED_OPERATIONS registry in pssh_sharder.py" + ) + self.fail(error_msg) + + print(f"✓ Sharder Compatibility Check Passed: All {len(used_operations)} operations supported") + + def test_sharder_operations_can_execute_without_errors(self): + """ + Test that all sharder operations can execute without parameter/method errors. + + Uses the operation registry to dynamically generate test payloads, + ensuring all operations are tested consistently. + """ + from cvs.lib.parallel.pssh_sharder import SUPPORTED_OPERATIONS + + config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) + sharder = PsshSharder(config) + + # Dynamically generate test payloads from operation registry + test_payloads = {} + + for operation_name, op_config in SUPPORTED_OPERATIONS.items(): + payload = { + 'operation': operation_name, + 'init': {'log': None, 'host_list': ['testhost'], 'user': 'test'}, + } + + # Add required parameters with test values + param_values = { + 'cmd': 'echo test', + 'cmd_list': ['echo test'], + 'timeout': 30, + 'detailed': False, + 'local_file': '/tmp/test.txt', + 'remote_file': '/remote/test.txt', + 'recurse': False, + 'suffix_separator': '_', + } + + for param in op_config['required_params']: + if param in param_values: + payload[param] = param_values[param] + + # Add some optional parameters for testing + for param in op_config['optional_params']: + if param in param_values: + payload[param] = param_values[param] + + test_payloads[operation_name] = payload + + # Mock all Pssh methods to avoid actual SSH calls + with patch('cvs.lib.parallel.pssh_sharder.Pssh') as mock_pssh_class: + mock_shard = MagicMock() + mock_pssh_class.return_value = mock_shard + + # Configure return values + mock_shard.exec.return_value = {'testhost': 'output'} + mock_shard.exec_cmd_list.return_value = {'testhost': 'output'} + mock_shard.download_file.return_value = {'testhost': '/tmp/test.txt_testhost'} + mock_shard.reachable_hosts = ['testhost'] + mock_shard.unreachable_hosts = [] + + failed_operations = [] + + for operation_name, payload in test_payloads.items(): + try: + result = sharder.run_shard(payload) + + # Verify result structure + self.assertIn('result', result) + self.assertIn('reachable_hosts', result) + self.assertIn('unreachable_hosts', result) + + # Verify correct method was called on mock shard + if operation_name == 'exec': + mock_shard.exec.assert_called() + elif operation_name == 'cmd_list': + mock_shard.exec_cmd_list.assert_called() + elif operation_name == 'scp': + mock_shard.scp_file.assert_called() + elif operation_name == 'upload': + mock_shard.upload_file.assert_called() + elif operation_name == 'download': + mock_shard.download_file.assert_called() + elif operation_name == 'reboot': + mock_shard.reboot_connections.assert_called() + + except Exception as e: + failed_operations.append( + {'operation': operation_name, 'error': str(e), 'error_type': type(e).__name__} + ) + finally: + # Reset mocks for next iteration + mock_shard.reset_mock() + + if failed_operations: + error_msg = "Sharder operations failed with runtime errors:\n" + for failure in failed_operations: + error_msg += f" {failure['operation']}: {failure['error_type']}: {failure['error']}\n" + error_msg += "\nThis suggests missing parameters, wrong method names, or incorrect implementation." + self.fail(error_msg) + + print(f"✓ Runtime Compatibility Check Passed: All {len(test_payloads)} operations execute correctly") + + def test_operation_registry_completeness(self): + """ + Test that the operation registry is complete and consistent. + + Validates that all operations have proper configuration and that + the referenced Pssh methods actually exist. + """ + from cvs.lib.parallel.pssh_sharder import SUPPORTED_OPERATIONS + from cvs.lib.parallel.pssh import Pssh + import inspect + + errors = [] + + for operation_name, op_config in SUPPORTED_OPERATIONS.items(): + # Check required fields + required_fields = ['required_params', 'optional_params', 'returns_result'] + for field in required_fields: + if field not in op_config: + errors.append(f"Operation '{operation_name}' missing required field '{field}'") + + # Check that operation name corresponds to an actual method in Pssh class + if not hasattr(Pssh, operation_name): + errors.append(f"Operation '{operation_name}' does not correspond to a method in Pssh class") + + # Check method signature compatibility + if hasattr(Pssh, operation_name): + try: + method = getattr(Pssh, operation_name) + sig = inspect.signature(method) + method_params = set(sig.parameters.keys()) - {'self'} # Exclude 'self' + + config_params = set(op_config['required_params'] + op_config['optional_params']) + + # Check if all config parameters exist in method signature + missing_in_method = config_params - method_params + if missing_in_method: + errors.append( + f"Operation '{operation_name}': method missing parameters {sorted(missing_in_method)}" + ) + + except Exception as e: + errors.append(f"Operation '{operation_name}': Could not validate method signature: {e}") + + if errors: + error_msg = "Operation registry validation failed:\n" + for error in errors: + error_msg += f" - {error}\n" + error_msg += "\nFix the SUPPORTED_OPERATIONS registry in pssh_sharder.py" + self.fail(error_msg) + + print(f"✓ Operation Registry Validation Passed: All {len(SUPPORTED_OPERATIONS)} operations properly configured") + + def test_registry_handler_consistency(self): + """ + Test that all operations in registry have corresponding handlers in run_shard. + + This prevents the implementation error scenario where operations are added to + the registry but corresponding handler code is forgotten. + """ + from cvs.lib.parallel.pssh_sharder import SUPPORTED_OPERATIONS + import inspect + import re + + # Get the run_shard method source to extract handled operations + from cvs.lib.parallel.pssh_sharder import PsshSharder + + run_shard_source = inspect.getsource(PsshSharder.run_shard) + + # Extract operation names from if/elif operation == 'name' and operation in ['name1', 'name2'] patterns + # Pattern 1: operation == 'name' + equals_pattern = r"operation == ['\"](\w+)['\"]" + equals_operations = set(re.findall(equals_pattern, run_shard_source)) + + # Pattern 2: operation in ['name1', 'name2', ...] + in_pattern = r"operation in \[(.*?)\]" + in_matches = re.findall(in_pattern, run_shard_source) + in_operations = set() + for match in in_matches: + # Extract individual operation names from the list + names = re.findall(r"['\"](\w+)['\"]", match) + in_operations.update(names) + + handled_operations = equals_operations | in_operations + + # Operations in registry but not in handlers (implementation error) + registry_operations = set(SUPPORTED_OPERATIONS.keys()) + missing_handlers = registry_operations - handled_operations + + # Operations in handlers but not in registry (less critical, but good to know) + extra_handlers = handled_operations - registry_operations + + errors = [] + if missing_handlers: + errors.append(f"Operations in registry but missing handlers: {sorted(missing_handlers)}") + + if extra_handlers: + errors.append(f"Operations with handlers but not in registry: {sorted(extra_handlers)}") + + if errors: + error_msg = "Registry-Handler consistency check failed:\n" + for error in errors: + error_msg += f" - {error}\n" + error_msg += "\nEnsure registry and handler code are synchronized." + self.fail(error_msg) + + print(f"✓ Registry-Handler Consistency Check Passed: All {len(registry_operations)} operations have handlers") + if __name__ == "__main__": unittest.main() From 7b58cb18f23c315a5b5ebbdc42033d93b725501f Mon Sep 17 00:00:00 2001 From: Ignatious Johnson Date: Tue, 12 May 2026 19:27:57 +0000 Subject: [PATCH 2/2] Refactor to ABC + composition + dynamic sharder architecture This commit addresses all PR review comments by completely refactoring the multiprocess SSH architecture to be cleaner, more maintainable, and more robust. Key Architectural Changes: - Replace inheritance with ABC + composition pattern - Eliminate registry + if/elif tree with dynamic method dispatch - Create single source of truth via ShardableSshInterface ABC - Move ABC to separate interfaces.py to avoid circular imports Review Comments Addressed: - #1-#2: Fixed operation name mismatches (upload/download) - #3: Updated test expectations for new calling convention - #4: Proper error handling with early ABC validation - #5: Maintained custom download_file result merging - #6: Implemented exact dynamic approach suggested by reviewer Co-authored-by: Cursor Signed-off-by: Ignatious Johnson --- cvs/lib/parallel/interfaces.py | 42 ++ cvs/lib/parallel/multiprocess_pssh.py | 105 ++-- cvs/lib/parallel/pssh.py | 37 -- cvs/lib/parallel/pssh_sharder.py | 120 +---- .../unittests/test_multiprocess_pssh.py | 477 +++++++----------- .../parallel/unittests/test_pssh_sharder.py | 334 ++---------- 6 files changed, 342 insertions(+), 773 deletions(-) create mode 100644 cvs/lib/parallel/interfaces.py diff --git a/cvs/lib/parallel/interfaces.py b/cvs/lib/parallel/interfaces.py new file mode 100644 index 00000000..f98c7e70 --- /dev/null +++ b/cvs/lib/parallel/interfaces.py @@ -0,0 +1,42 @@ +''' +Copyright 2025 Advanced Micro Devices, Inc. +All rights reserved. This notice is intended as a precaution against inadvertent publication and does not imply publication or any waiver of confidentiality. +The year included in the foregoing notice is the year of creation of the work. +All code contained here is Property of Advanced Micro Devices, Inc. +''' + +from abc import ABC, abstractmethod + + +class ShardableSshInterface(ABC): + """Abstract base class defining operations that MUST support sharding. + + Any class implementing this interface must provide sharded implementations + for all these methods. This prevents silent performance degradation where + operations fall back to non-sharded implementations. + """ + + @abstractmethod + def exec(self, cmd, timeout=None, print_console=True, detailed=False): + """Execute command - must support sharding for performance.""" + pass + + @abstractmethod + def exec_cmd_list(self, cmd_list, timeout=None, print_console=True): + """Execute command list - must support sharding for performance.""" + pass + + @abstractmethod + def upload_file(self, local_file, remote_file, recurse=False): + """Upload file via SFTP - must support sharding for performance.""" + pass + + @abstractmethod + def download_file(self, remote_file, local_file, recurse=False, suffix_separator='_'): + """Download file via SFTP - must support sharding for performance.""" + pass + + @abstractmethod + def reboot_connections(self): + """Reboot connections - must support sharding for performance.""" + pass diff --git a/cvs/lib/parallel/multiprocess_pssh.py b/cvs/lib/parallel/multiprocess_pssh.py index 7b47ffcd..4d9ce9fb 100644 --- a/cvs/lib/parallel/multiprocess_pssh.py +++ b/cvs/lib/parallel/multiprocess_pssh.py @@ -9,12 +9,13 @@ from cvs.lib.parallel.pssh import Pssh from cvs.lib.parallel.config import ParallelConfig from cvs.lib.parallel.pssh_sharder import PsshSharder +from cvs.lib.parallel.interfaces import ShardableSshInterface from cvs.lib import globals global_log = globals.log -class MultiProcessPssh(Pssh): +class MultiProcessPssh(ShardableSshInterface): """ Multi-process parallel SSH with automatic host sharding for large host lists. @@ -47,10 +48,16 @@ def __init__( n = len(host_list) if host_list is not None else 0 use_mp = hosts_per_shard > 0 and n > hosts_per_shard - if not use_mp: - # Use single-process base class - super().__init__( - log, # Pass through log parameter to parent + if use_mp: + # Initialize for multi-process sharding - no Pssh instance needed + self._init_sharded(log, host_list, user, password, pkey, host_key_check, stop_on_errors, env_vars) + self.sharder = PsshSharder(self.config) + self._use_process_sharding = True + self.pssh = None # No Pssh instance needed for sharded mode + else: + # No sharding - create Pssh instance for delegation + self.pssh = Pssh( + log, host_list, user, password, @@ -63,12 +70,6 @@ def __init__( self._use_process_sharding = False # Ensure attributes needed by _shard_init_kwargs are available self.env_vars = env_vars - else: - # Initialize for multi-process sharding - self._init_sharded(log, host_list, user, password, pkey, host_key_check, stop_on_errors, env_vars) - - # Create the sharder - no registration needed with direct operations! - self.sharder = PsshSharder(self.config) def _init_sharded( self, @@ -116,9 +117,8 @@ def _shard_init_kwargs(self): def _merge_shard_returns(self, shard_returns, merge_unreachable=True): """ - Update reachable/unreachable host lists from shard workers and convert results to cmd_output dict. + Update reachable/unreachable host lists from shard workers and merge results to cmd_output dict. - Handles both state updates and SimpleHostOutput conversion in one place for cleaner architecture. Returns cmd_output dict ready for _print_merged_outputs. """ # 1. Update host lists (existing logic) @@ -132,18 +132,12 @@ def _merge_shard_returns(self, shard_returns, merge_unreachable=True): if u not in self.unreachable_hosts: self.unreachable_hosts.append(u) - # 2. Convert SimpleHostOutput objects to cmd_output dict + # 2. Merge results from shard workers (always dict format) merged = {} for shard in shard_returns: result = shard.get('result') or {} # treat None as {} (for scp/reboot) - - if isinstance(result, list): - # Handle SimpleHostOutput objects - convert using _process_output - temp_dict = self._process_output(result, print_console=False) - merged.update(temp_dict) - else: - # Handle old dict format (for scp/reboot) - merged.update(result) + # Shard workers always return processed dicts from _process_output + merged.update(result) # 3. Preserve original host order cmd_output = {} @@ -173,8 +167,8 @@ def _print_merged_outputs(self, cmd_output, cmd=None, cmd_list=None, print_conso def exec(self, cmd, timeout=None, print_console=True, detailed=False): """Execute command with automatic sharding if needed.""" - if not getattr(self, '_use_process_sharding', False): - return super().exec(cmd, timeout=timeout, print_console=print_console, detailed=detailed) + if not self._use_process_sharding: + return self.pssh.exec(cmd, timeout=timeout, print_console=print_console, detailed=detailed) if self.env_prefix: full_cmd = f"{self.env_prefix} ; {cmd}" @@ -213,8 +207,8 @@ def exec(self, cmd, timeout=None, print_console=True, detailed=False): def exec_cmd_list(self, cmd_list, timeout=None, print_console=True): """Execute command list with automatic sharding if needed.""" - if not getattr(self, '_use_process_sharding', False): - return super().exec_cmd_list(cmd_list, timeout=timeout, print_console=print_console) + if not self._use_process_sharding: + return self.pssh.exec_cmd_list(cmd_list, timeout=timeout, print_console=print_console) if len(cmd_list) != len(self.host_list): raise ValueError('cmd_list length must match host_list length') @@ -275,28 +269,20 @@ def exec_cmd_list(self, cmd_list, timeout=None, print_console=True): return cmd_output def scp_file(self, local_file, remote_file, recurse=False): - """Copy file with automatic sharding if needed.""" - if not getattr(self, '_use_process_sharding', False): - return super().scp_file(local_file, remote_file, recurse=recurse) + """ + Backward-compatible alias for upload_file. + Kept so existing callers (and log-grep tooling looking for the legacy + "About to copy local file..." line) keep working. New code should call + upload_file directly. + """ self.log.info('About to copy local file {} to remote {} on all Hosts'.format(local_file, remote_file)) - - host_chunks = list(self.sharder.chunk_hosts(self.reachable_hosts)) - payloads = self.sharder.create_payloads( - 'scp_file', - host_chunks, - self._shard_init_kwargs(), - local_file=local_file, - remote_file=remote_file, - recurse=recurse, - ) - shard_returns = self.sharder.execute_sharded(payloads) - return self._merge_shard_returns(shard_returns, merge_unreachable=False) + return self.upload_file(local_file, remote_file, recurse=recurse) def upload_file(self, local_file, remote_file, recurse=False): """Upload file with automatic sharding if needed.""" - if not getattr(self, '_use_process_sharding', False): - return super().upload_file(local_file, remote_file, recurse=recurse) + if not self._use_process_sharding: + return self.pssh.upload_file(local_file, remote_file, recurse=recurse) self.log.info('SFTP upload %s -> %s on %s', local_file, remote_file, self.reachable_hosts) @@ -310,12 +296,12 @@ def upload_file(self, local_file, remote_file, recurse=False): recurse=recurse, ) shard_returns = self.sharder.execute_sharded(payloads) - return self._merge_shard_returns(shard_returns, merge_unreachable=False) + return self._merge_shard_returns(shard_returns) def download_file(self, remote_file, local_file, recurse=False, suffix_separator='_'): """Download file with automatic sharding if needed.""" - if not getattr(self, '_use_process_sharding', False): - return super().download_file(remote_file, local_file, recurse=recurse, suffix_separator=suffix_separator) + if not self._use_process_sharding: + return self.pssh.download_file(remote_file, local_file, recurse=recurse, suffix_separator=suffix_separator) self.log.info('SFTP download %s -> %s from %s', remote_file, local_file, self.reachable_hosts) @@ -330,21 +316,12 @@ def download_file(self, remote_file, local_file, recurse=False, suffix_separator suffix_separator=suffix_separator, ) shard_returns = self.sharder.execute_sharded(payloads) - - # For download_file, we need to merge the returned dictionaries from each shard - # The result is a dict mapping {host: actual_local_path} - merged_result = {} - for shard in shard_returns: - result = shard.get('result') or {} - if isinstance(result, dict): - merged_result.update(result) - - return merged_result + return self._merge_shard_returns(shard_returns) def reboot_connections(self): """Reboot connections with automatic sharding if needed.""" - if not getattr(self, '_use_process_sharding', False): - return super().reboot_connections() + if not self._use_process_sharding: + return self.pssh.reboot_connections() self.log.info('Rebooting Connections') @@ -355,8 +332,12 @@ def reboot_connections(self): def destroy_clients(self): """Destroy clients - handle both sharded and non-sharded modes.""" - if getattr(self, '_use_process_sharding', False): - self.log.info('Destroying Current phdl connections ..') + if self._use_process_sharding: + self.log.info('Cleaning up sharded mode state ..') + # In sharded mode, no persistent connections to destroy self.client = None - return - super().destroy_clients() + else: + self.log.info('Destroying Current phdl connections ..') + # In non-sharded mode, properly destroy the Pssh instance connections + if self.pssh is not None: + self.pssh.destroy_clients() diff --git a/cvs/lib/parallel/pssh.py b/cvs/lib/parallel/pssh.py index 1354dd75..26a58203 100644 --- a/cvs/lib/parallel/pssh.py +++ b/cvs/lib/parallel/pssh.py @@ -17,17 +17,6 @@ global_log = globals.log -class SimpleHostOutput: - """Simple HostOutput-compatible object for consistent processing across process boundaries.""" - - def __init__(self, host, stdout_lines, stderr_lines, exception, exit_code=0): - self.host = host - self.stdout = iter(stdout_lines) # Convert list to iterator for _process_output compatibility - self.stderr = iter(stderr_lines) # Convert list to iterator for _process_output compatibility - self.exception = exception - self.exit_code = exit_code - - class Pssh: """ Single-process parallel SSH: one ParallelSSHClient (one gevent hub) over a host list. @@ -385,32 +374,6 @@ def reboot_connections(self): self.log.info('Rebooting Connections') self.client.run_command('reboot -f', stop_on_errors=self.stop_on_errors) - def _extract_simple_data(self, output): - """ - Extract essential data from pssh.output.HostOutput objects for safe IPC transfer. - - Converts non-picklable HostOutput objects (containing active SSH channels, - generators, and C extension objects) into picklable SimpleHostOutput objects - by consuming stdout/stderr generators into lists and extracting core attributes. - """ - - simple_outputs = [] - for item in output: - # Consume generators to lists for pickling - stdout_lines = list(item.stdout or []) - stderr_lines = list(item.stderr or []) - - simple_output = SimpleHostOutput( - host=item.host, - stdout_lines=stdout_lines, - stderr_lines=stderr_lines, - exception=item.exception, - exit_code=getattr(item, 'exit_code', 0), - ) - simple_outputs.append(simple_output) - - return simple_outputs - def destroy_clients(self): self.log.info('Destroying Current phdl connections ..') del self.client diff --git a/cvs/lib/parallel/pssh_sharder.py b/cvs/lib/parallel/pssh_sharder.py index f15c9b19..3c3d47dc 100644 --- a/cvs/lib/parallel/pssh_sharder.py +++ b/cvs/lib/parallel/pssh_sharder.py @@ -5,45 +5,18 @@ All code contained here is Property of Advanced Micro Devices, Inc. ''' +import inspect import multiprocessing as mp from concurrent.futures import ProcessPoolExecutor, as_completed from cvs.lib.parallel.pssh import Pssh +from cvs.lib.parallel.interfaces import ShardableSshInterface -# Operation registry - single source of truth for supported operations -# Operation names now match method names for consistency and simplicity +# Dynamically discover supported operations from ABC (computed once at import time) + SUPPORTED_OPERATIONS = { - 'exec': { - 'required_params': ['cmd'], - 'optional_params': ['timeout', 'detailed'], - 'returns_result': True, - }, - 'exec_cmd_list': { - 'required_params': ['cmd_list'], - 'optional_params': ['timeout'], - 'returns_result': True, - }, - 'scp_file': { - 'required_params': ['local_file', 'remote_file'], - 'optional_params': ['recurse'], - 'returns_result': False, - }, - 'upload_file': { - 'required_params': ['local_file', 'remote_file'], - 'optional_params': ['recurse'], - 'returns_result': False, - }, - 'download_file': { - 'required_params': ['remote_file', 'local_file'], - 'optional_params': ['recurse', 'suffix_separator'], - 'returns_result': True, - }, - 'reboot_connections': { - 'required_params': [], - 'optional_params': [], - 'returns_result': False, - }, + name for name, method in inspect.getmembers(ShardableSshInterface) if getattr(method, '__isabstractmethod__', False) } @@ -110,6 +83,8 @@ def run_shard(payload): """ Run an SSH operation on a shard of hosts (must be picklable for multiprocessing). + Dynamically supports all abstract methods defined in ShardableSshInterface. + Args: payload: Dict with 'operation' (operation type), 'init' (SSH setup), and operation args @@ -120,80 +95,35 @@ def run_shard(payload): if not isinstance(operation, str): raise TypeError('payload["operation"] must be str, got %r' % (type(operation),)) + # Validate operation is supported by ABC + if operation not in SUPPORTED_OPERATIONS: + raise ValueError(f'Unknown operation: {operation}. Supported: {sorted(SUPPORTED_OPERATIONS)}') + # Create SSH client for this shard of hosts init_kwargs = payload['init'] init_kwargs['process_output'] = False # Force raw output mode for sharding shard = Pssh(**init_kwargs) try: - # Use operation registry for parameter validation and documentation - op_config = SUPPORTED_OPERATIONS.get(operation) - if not op_config: - # Let the else clause in the operation handler catch unknown operations - op_config = {'required_params': [], 'optional_params': [], 'returns_result': False} - method_name = operation # Operation name IS the method name! - shard_method = getattr(shard, method_name) - - # Build arguments from payload based on operation configuration - args = {} - - # Add required parameters - for param in op_config['required_params']: - if param not in payload: - raise ValueError(f"Missing required parameter '{param}' for operation '{operation}'") - args[param] = payload[param] - - # Add optional parameters if present - for param in op_config['optional_params']: - if param in payload: - args[param] = payload[param] + # Ensure method exists in Pssh + if not hasattr(shard, operation): + raise RuntimeError(f'Method {operation} not found in Pssh class') + + shard_method = getattr(shard, operation) + + # Extract operation arguments from payload (excluding 'operation' and 'init') + args = {k: v for k, v in payload.items() if k not in ['operation', 'init']} # Add common parameters for operations that support them - if method_name in ['exec', 'exec_cmd_list']: + if operation in ['exec', 'exec_cmd_list']: args['print_console'] = False # Always False for sharded operations - # Execute the operation - handle special cases for method signatures - if operation == 'exec': - # Maintain existing call signature for backward compatibility - kwargs = { - 'timeout': args.get('timeout'), - 'print_console': False, - } - if 'detailed' in args: - kwargs['detailed'] = args['detailed'] - result = shard_method(args['cmd'], **kwargs) - elif operation == 'exec_cmd_list': - # Maintain existing call signature for backward compatibility - result = shard_method(args['cmd_list'], timeout=args.get('timeout'), print_console=False) - elif operation in ['scp_file', 'upload_file']: - # File upload operations - shard_method(args['local_file'], args['remote_file'], recurse=args.get('recurse', False)) - result = None - elif operation == 'download_file': - # File download operation - result = shard_method( - args['remote_file'], - args['local_file'], - recurse=args.get('recurse', False), - suffix_separator=args.get('suffix_separator', '_'), - ) - elif operation == 'reboot_connections': - # Reboot operation - shard_method() + # Call the method dynamically + result = shard_method(**args) + + # Handle operations that should return None (void operations) + if operation in ['upload_file', 'reboot_connections']: result = None - else: - # Enhanced error handling to distinguish user vs implementation errors - if operation in SUPPORTED_OPERATIONS: - # Case #2: Implementation Error - Operation in registry but no handler - raise RuntimeError( - f'Implementation bug: Operation \'{operation}\' is supported (in registry) ' - f'but handler is missing in run_shard().' - ) - else: - # Case #1: User Error - Operation not supported at all - raise ValueError( - f'Unknown operation: {operation}. Supported operations: {list(SUPPORTED_OPERATIONS.keys())}' - ) return { 'result': result, diff --git a/cvs/lib/parallel/unittests/test_multiprocess_pssh.py b/cvs/lib/parallel/unittests/test_multiprocess_pssh.py index 6896a07f..a6c251d4 100644 --- a/cvs/lib/parallel/unittests/test_multiprocess_pssh.py +++ b/cvs/lib/parallel/unittests/test_multiprocess_pssh.py @@ -8,22 +8,30 @@ class TestMultiProcessPsshInitialization(unittest.TestCase): def setUp(self): self.mock_log = MagicMock() self.host_list = ["host1", "host2", "host3", "host4", "host5"] + # Mock Pssh.__init__ for all tests in this class + self.pssh_patcher = patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') + self.mock_pssh_init = self.pssh_patcher.start() + self.mock_pssh_init.return_value = None - @patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') - def test_init_no_sharding_small_host_list(self, mock_parent_init): + def tearDown(self): + self.pssh_patcher.stop() + + def test_init_no_sharding_small_host_list(self): """Test initialization without sharding for small host lists.""" - mock_parent_init.return_value = None small_host_list = ["host1", "host2"] config = ParallelConfig(hosts_per_shard=32) pssh = MultiProcessPssh(self.mock_log, small_host_list, user="test", config=config) - # Should call parent __init__ for small lists - mock_parent_init.assert_called_once_with( + # Always creates composed Pssh instance with ABC+composition + self.mock_pssh_init.assert_called_once_with( self.mock_log, small_host_list, "test", None, 'id_rsa', False, True, None, process_output=True ) - # For small lists, no sharder should be created + # For small lists, no sharder should be created (key difference!) self.assertFalse(hasattr(pssh, 'sharder')) + self.assertFalse(pssh._use_process_sharding) + # But pssh instance should always exist + self.assertTrue(hasattr(pssh, 'pssh')) @patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded') @patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') @@ -33,174 +41,120 @@ def test_init_with_sharding_large_host_list(self, mock_sharder_class, mock_init_ mock_sharder = MagicMock() mock_sharder_class.return_value = mock_sharder - MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + # In sharded mode, should NOT create Pssh instance (avoids bottleneck) + self.mock_pssh_init.assert_not_called() # Should call _init_sharded for large lists mock_init_sharded.assert_called_once_with( self.mock_log, self.host_list, "test", None, 'id_rsa', False, True, None ) # Verify sharder was created mock_sharder_class.assert_called_once_with(config) + self.assertTrue(pssh._use_process_sharding) + # In sharded mode: pssh is None, sharder exists + self.assertIsNone(pssh.pssh) + self.assertTrue(hasattr(pssh, 'sharder')) def test_init_with_custom_config(self): """Test initialization with custom configuration.""" config = ParallelConfig(hosts_per_shard=64, max_workers_per_cpu=8) - with patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') as mock_parent_init: - mock_parent_init.return_value = None - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", password="pass", config=config) - - self.assertEqual(pssh.config, config) - mock_parent_init.assert_called_once() - - def test_init_default_config(self): - """Test initialization with default configuration.""" - with patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') as mock_parent_init: - mock_parent_init.return_value = None - pssh = MultiProcessPssh(self.mock_log, ["host1"], user="test") - - # Should create default config - self.assertIsNotNone(pssh.config) - self.assertIsInstance(pssh.config, ParallelConfig) - - -class TestMultiProcessPsshInitSharded(unittest.TestCase): - def setUp(self): - self.mock_log = MagicMock() - self.host_list = ["host1", "host2"] - - @patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') - def test_init_sharded_sets_attributes(self, mock_sharder_class): - """Test that _init_sharded sets all required attributes.""" - pssh = MultiProcessPssh.__new__(MultiProcessPssh) # Create without calling __init__ - # Set up config and sharder that _init_sharded expects - pssh.config = ParallelConfig(hosts_per_shard=2, max_workers_per_cpu=1) - mock_sharder = MagicMock() - pssh.sharder = mock_sharder - - pssh._init_sharded( - self.mock_log, self.host_list, "testuser", "testpass", "custom_key", True, False, {"VAR": "value"} - ) + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", password="pass", config=config) - # Logger is now module-level, no longer stored in pssh.log - self.assertEqual(pssh.host_list, self.host_list) - self.assertEqual(pssh.reachable_hosts, self.host_list) - self.assertEqual(pssh.user, "testuser") - self.assertEqual(pssh.password, "testpass") - self.assertEqual(pssh.pkey, "custom_key") - self.assertTrue(pssh.host_key_check) - self.assertFalse(pssh.stop_on_errors) - self.assertEqual(pssh.env_vars, {"VAR": "value"}) - self.assertTrue(pssh._use_process_sharding) + self.assertEqual(pssh.config, config) class TestMultiProcessPsshExec(unittest.TestCase): def setUp(self): self.mock_log = MagicMock() self.host_list = ["host1", "host2"] + # Mock Pssh.__init__ for all tests in this class + self.pssh_patcher = patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') + self.mock_pssh_init = self.pssh_patcher.start() + self.mock_pssh_init.return_value = None - def test_exec_no_sharding(self): - """Test exec without sharding calls parent method.""" - with patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') as mock_parent_init: - with patch('cvs.lib.parallel.multiprocess_pssh.Pssh.exec') as mock_parent_exec: - mock_parent_init.return_value = None - mock_parent_exec.return_value = {"host1": "result"} - - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test") - result = pssh.exec("uptime") + # Mock execute_sharded to avoid process spawning while testing real payload creation + self.execute_sharded_patcher = patch('cvs.lib.parallel.pssh_sharder.PsshSharder.execute_sharded') + self.mock_execute_sharded = self.execute_sharded_patcher.start() - mock_parent_exec.assert_called_once_with("uptime", timeout=None, print_console=True, detailed=False) - self.assertEqual(result, {"host1": "result"}) + def tearDown(self): + self.pssh_patcher.stop() + self.execute_sharded_patcher.stop() def test_exec_with_sharding(self): """Test exec with sharding uses sharder.""" config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) - with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): - with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: - mock_sharder = MagicMock() - mock_sharder_class.return_value = mock_sharder + # Mock execute_sharded to return realistic shard results + self.mock_execute_sharded.return_value = [ + {'result': {"host1": "up1"}, 'reachable_hosts': ["host1"], 'unreachable_hosts': []}, + {'result': {"host2": "up2"}, 'reachable_hosts': ["host2"], 'unreachable_hosts': []}, + ] - # Set up the mock sharder behavior - mock_sharder.chunk_hosts.return_value = [["host1"], ["host2"]] - mock_sharder.create_payloads.return_value = [ - {"operation": "exec", "init": {}, "cmd": "uptime"}, - {"operation": "exec", "init": {}, "cmd": "uptime"}, - ] - mock_sharder.execute_sharded.return_value = [ - {"result": {"host1": "up1"}, "reachable_hosts": ["host1"], "unreachable_hosts": []}, - {"result": {"host2": "up2"}, "reachable_hosts": ["host2"], "unreachable_hosts": []}, - ] - mock_sharder.merge_results.return_value = {"host1": "up1", "host2": "up2"} + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.exec("uptime", timeout=30, detailed=True) - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + # Verify execute_sharded was called with real payloads from create_payloads + self.mock_execute_sharded.assert_called_once() + payloads = self.mock_execute_sharded.call_args[0][0] - # Manually set attributes that _init_sharded would set - pssh.env_prefix = None - pssh.host_list = self.host_list - pssh.reachable_hosts = self.host_list - pssh.sharder = mock_sharder - pssh._use_process_sharding = True # Enable sharding for this test - # Add attributes needed by _shard_init_kwargs() - pssh.user = "test" - pssh.password = None - pssh.pkey = "id_rsa" - pssh.host_key_check = False - pssh.stop_on_errors = True - pssh.env_vars = None + # Check that real payloads were created correctly + self.assertEqual(len(payloads), 2) # Two shards (hosts_per_shard=1) + first_payload = payloads[0] + self.assertEqual(first_payload['operation'], 'exec') + self.assertEqual(first_payload['cmd'], 'uptime') + self.assertEqual(first_payload['timeout'], 30) + self.assertEqual(first_payload['detailed'], True) + self.assertIn('init', first_payload) - result = pssh.exec("uptime") + # Verify real init kwargs are properly constructed by _shard_init_kwargs + init_kwargs = first_payload['init'] + self.assertEqual(init_kwargs['user'], 'test') + self.assertIn('host_list', init_kwargs) - # Verify sharder methods were called - mock_sharder.chunk_hosts.assert_called_once_with(self.host_list) - mock_sharder.execute_sharded.assert_called_once() - self.assertEqual(result, {"host1": "up1", "host2": "up2"}) + # Verify result is merged from shards + self.assertEqual(result, {"host1": "up1", "host2": "up2"}) def test_exec_cmd_list_with_sharding(self): """Test exec_cmd_list with sharding uses sharder.""" config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) cmd_list = ["uptime", "date"] - with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): - with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: - mock_sharder = MagicMock() - mock_sharder_class.return_value = mock_sharder + # Mock execute_sharded to capture real payloads and return results + def capture_payloads_and_return_results(payloads): + self.captured_payloads = payloads + return [ + { + 'result': {"host1": "up1", "host2": "date2"}, + 'reachable_hosts': ["host1", "host2"], + 'unreachable_hosts': [], + } + ] - # Set up the mock sharder behavior - create_payloads returns a list, we need the [0] element - mock_sharder.create_payloads.return_value = [ - {"operation": "cmd_list", "init": {}, "cmd_list": cmd_list} - ] - mock_sharder.execute_sharded.return_value = [ - { - "result": {"host1": "up1", "host2": "date2"}, - "reachable_hosts": ["host1", "host2"], - "unreachable_hosts": [], - }, - ] - mock_sharder.merge_results.return_value = {"host1": "up1", "host2": "date2"} + self.mock_execute_sharded.side_effect = capture_payloads_and_return_results - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.exec_cmd_list(cmd_list, timeout=60) - # Manually set attributes that _init_sharded would set - pssh.env_prefix = None - pssh.host_list = self.host_list - pssh.reachable_hosts = self.host_list - pssh.config = config - pssh.sharder = mock_sharder - pssh._use_process_sharding = True # Enable sharding for this test - # Additional attributes needed by _shard_init_kwargs - pssh.user = "test" - pssh.password = None - pssh.pkey = "id_rsa" - pssh.host_key_check = False - pssh.stop_on_errors = True - pssh.env_vars = None + # Verify execute_sharded was called with real payloads + self.mock_execute_sharded.assert_called_once() - result = pssh.exec_cmd_list(cmd_list) + # Check that real payloads were created correctly by create_payloads + self.assertEqual(len(self.captured_payloads), 2) # Two shards + first_payload = self.captured_payloads[0] + self.assertEqual(first_payload['operation'], 'exec_cmd_list') + # With hosts_per_shard=1, each shard gets one command from the list + self.assertEqual(first_payload['cmd_list'], ['uptime']) # First shard gets first command + self.assertEqual(first_payload['timeout'], 60) + self.assertIn('init', first_payload) + + # Verify real init kwargs + init_kwargs = first_payload['init'] + self.assertEqual(init_kwargs['user'], 'test') - # Verify sharder methods were called - mock_sharder.execute_sharded.assert_called_once() - self.assertEqual(result, {"host1": "up1", "host2": "date2"}) + # Verify result is merged from shards + self.assertEqual(result, {"host1": "up1", "host2": "date2"}) def test_exec_cmd_list_with_unreachable_hosts(self): """Test exec_cmd_list correctly maps commands when some hosts are unreachable.""" @@ -219,7 +173,11 @@ def test_exec_cmd_list_with_unreachable_hosts(self): # Mock sharder behavior mock_sharder.chunk_hosts.return_value = [["host1", "host3"], ["host4"]] mock_sharder.create_payloads.return_value = [ - {"operation": "cmd_list", "init": {}, "cmd_list": ["cmd1", "cmd3"]}, # Expected: correct commands + { + "operation": "exec_cmd_list", + "init": {}, + "cmd_list": ["cmd1", "cmd3"], + }, # Expected: correct commands ] mock_sharder.execute_sharded.return_value = [ { @@ -260,98 +218,76 @@ def test_exec_cmd_list_with_unreachable_hosts(self): self.assertEqual(result, {"host1": "result1", "host3": "result3", "host4": "result4"}) def test_scp_file_with_sharding(self): - """Test scp_file with sharding uses sharder.""" + """Test scp_file delegates to upload_file with sharding.""" config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) - with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): - with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: - mock_sharder = MagicMock() - mock_sharder_class.return_value = mock_sharder + # Mock execute_sharded for void scp_file operation + self.mock_execute_sharded.return_value = [ + {'result': None, 'reachable_hosts': ["host1"], 'unreachable_hosts': []}, + {'result': None, 'reachable_hosts': ["host2"], 'unreachable_hosts': []}, + ] - # Set up the mock sharder behavior - mock_sharder.chunk_hosts.return_value = [["host1"], ["host2"]] - mock_sharder.create_payloads.return_value = [ - {"operation": "scp", "init": {}, "local_file": "test.txt", "remote_file": "/tmp/test.txt"}, - ] - mock_sharder.execute_sharded.return_value = [ - { - "result": {"host1": "scp1", "host2": "scp2"}, - "reachable_hosts": ["host1", "host2"], - "unreachable_hosts": [], - }, - ] - mock_sharder.merge_results.return_value = {"host1": "scp1", "host2": "scp2"} + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.scp_file("test.txt", "/tmp/test.txt", recurse=False) - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) - - # Manually set attributes that _init_sharded would set - pssh.host_list = self.host_list - pssh.reachable_hosts = self.host_list - pssh.sharder = mock_sharder - pssh._use_process_sharding = True # Enable sharding for this test - # Add attributes needed by _shard_init_kwargs() - pssh.user = "test" - pssh.password = None - pssh.pkey = "id_rsa" - pssh.host_key_check = False - pssh.stop_on_errors = True - pssh.env_vars = None + # Verify execute_sharded was called with real payloads + self.mock_execute_sharded.assert_called_once() + payloads = self.mock_execute_sharded.call_args[0][0] - result = pssh.scp_file("test.txt", "/tmp/test.txt") + # Check that real payloads were created for upload_file operation + self.assertEqual(len(payloads), 2) # Two shards + first_payload = payloads[0] + self.assertEqual(first_payload['operation'], 'upload_file') # scp_file delegates to upload_file + self.assertEqual(first_payload['local_file'], 'test.txt') + self.assertEqual(first_payload['remote_file'], '/tmp/test.txt') + self.assertEqual(first_payload['recurse'], False) - # Verify sharder methods were called - mock_sharder.chunk_hosts.assert_called_once_with(self.host_list) - mock_sharder.execute_sharded.assert_called_once() - self.assertEqual(result, {"host1": "scp1", "host2": "scp2"}) + # scp_file should return empty dict (merged void results) + self.assertEqual(result, {}) def test_reboot_connections_with_sharding(self): """Test reboot_connections with sharding uses sharder.""" config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) - with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): - with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: - mock_sharder = MagicMock() - mock_sharder_class.return_value = mock_sharder + # Mock execute_sharded for void reboot_connections operation + self.mock_execute_sharded.return_value = [ + {'result': None, 'reachable_hosts': ["host1"], 'unreachable_hosts': []}, + {'result': None, 'reachable_hosts': ["host2"], 'unreachable_hosts': []}, + ] - # Set up the mock sharder behavior - mock_sharder.chunk_hosts.return_value = [["host1"], ["host2"]] - mock_sharder.create_payloads.return_value = [{"operation": "reboot", "init": {}}] - mock_sharder.execute_sharded.return_value = [ - { - "result": {"host1": "reboot1", "host2": "reboot2"}, - "reachable_hosts": ["host1", "host2"], - "unreachable_hosts": [], - }, - ] - mock_sharder.merge_results.return_value = {"host1": "reboot1", "host2": "reboot2"} + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.reboot_connections() - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + # Verify execute_sharded was called with real payloads + self.mock_execute_sharded.assert_called_once() + payloads = self.mock_execute_sharded.call_args[0][0] - # Manually set attributes that _init_sharded would set - pssh.host_list = self.host_list - pssh.reachable_hosts = self.host_list - pssh.sharder = mock_sharder - pssh._use_process_sharding = True # Enable sharding for this test - # Add attributes needed by _shard_init_kwargs() - pssh.user = "test" - pssh.password = None - pssh.pkey = "id_rsa" - pssh.host_key_check = False - pssh.stop_on_errors = True - pssh.env_vars = None + # Check that real payloads were created correctly + self.assertEqual(len(payloads), 2) # Two shards + first_payload = payloads[0] + self.assertEqual(first_payload['operation'], 'reboot_connections') + self.assertIn('init', first_payload) - result = pssh.reboot_connections() - - # Verify sharder methods were called - mock_sharder.chunk_hosts.assert_called_once_with(self.host_list) - mock_sharder.execute_sharded.assert_called_once() - self.assertEqual(result, {"host1": "reboot1", "host2": "reboot2"}) + # reboot_connections should return empty dict (merged void results) + self.assertEqual(result, {}) class TestMultiProcessPsshHelperMethods(unittest.TestCase): def setUp(self): self.mock_log = MagicMock() self.host_list = ["host1", "host2"] + # Mock Pssh.__init__ for all tests in this class + self.pssh_patcher = patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') + self.mock_pssh_init = self.pssh_patcher.start() + self.mock_pssh_init.return_value = None + + # Mock execute_sharded to avoid process spawning while testing real payload creation + self.execute_sharded_patcher = patch('cvs.lib.parallel.pssh_sharder.PsshSharder.execute_sharded') + self.mock_execute_sharded = self.execute_sharded_patcher.start() + + def tearDown(self): + self.pssh_patcher.stop() + self.execute_sharded_patcher.stop() def test_shard_init_kwargs(self): """Test _shard_init_kwargs creates correct initialization arguments.""" @@ -409,100 +345,57 @@ def test_upload_file_with_sharding(self): """Test upload_file with sharding uses sharder.""" config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) - with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): - with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: - mock_sharder = MagicMock() - mock_sharder_class.return_value = mock_sharder - - # Set up the mock sharder behavior - mock_sharder.chunk_hosts.return_value = [["host1"], ["host2"]] - mock_sharder.create_payloads.return_value = [ - {"operation": "upload", "init": {}, "local_file": "test.txt", "remote_file": "/tmp/test.txt"}, - ] - mock_sharder.execute_sharded.return_value = [ - { - "result": None, # upload_file returns None on success - "reachable_hosts": ["host1", "host2"], - "unreachable_hosts": [], - }, - ] + # Mock execute_sharded for void upload_file operation + self.mock_execute_sharded.return_value = [ + {'result': None, 'reachable_hosts': ["host1"], 'unreachable_hosts': []}, + {'result': None, 'reachable_hosts': ["host2"], 'unreachable_hosts': []}, + ] - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.upload_file("test.txt", "/tmp/test.txt", recurse=True) - # Manually set attributes that _init_sharded would set - pssh.host_list = self.host_list - pssh.reachable_hosts = self.host_list - pssh.sharder = mock_sharder - pssh._use_process_sharding = True # Enable sharding for this test - # Add attributes needed by _shard_init_kwargs() - pssh.user = "test" - pssh.password = None - pssh.pkey = "id_rsa" - pssh.host_key_check = False - pssh.stop_on_errors = True - pssh.env_vars = None + # Verify execute_sharded was called with real payloads + self.mock_execute_sharded.assert_called_once() + payloads = self.mock_execute_sharded.call_args[0][0] - result = pssh.upload_file("test.txt", "/tmp/test.txt", recurse=True) + # Check that real payload was created correctly + self.assertEqual(len(payloads), 2) # Two shards + first_payload = payloads[0] + self.assertEqual(first_payload['operation'], 'upload_file') + self.assertEqual(first_payload['local_file'], 'test.txt') + self.assertEqual(first_payload['remote_file'], '/tmp/test.txt') + self.assertEqual(first_payload['recurse'], True) + self.assertIn('init', first_payload) - # Verify sharder methods were called - mock_sharder.chunk_hosts.assert_called_once_with(self.host_list) - mock_sharder.execute_sharded.assert_called_once() - # upload_file should return the merged results (empty dict for success) - self.assertEqual(result, {}) + # upload_file should return empty dict (merged void results) + self.assertEqual(result, {}) def test_download_file_with_sharding(self): """Test download_file with sharding uses sharder and merges results.""" config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) - with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): - with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: - mock_sharder = MagicMock() - mock_sharder_class.return_value = mock_sharder - - # Set up the mock sharder behavior - mock_sharder.chunk_hosts.return_value = [["host1"], ["host2"]] - mock_sharder.create_payloads.return_value = [ - { - "operation": "download", - "init": {}, - "remote_file": "/remote/test.txt", - "local_file": "/tmp/test.txt", - }, - ] - # Simulate download results from each shard - mock_sharder.execute_sharded.return_value = [ - { - "result": {"host1": "/tmp/test.txt_host1"}, # Download results from shard 1 - "reachable_hosts": ["host1"], - "unreachable_hosts": [], - }, - { - "result": {"host2": "/tmp/test.txt_host2"}, # Download results from shard 2 - "reachable_hosts": ["host2"], - "unreachable_hosts": [], - }, - ] - - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) - - # Manually set attributes that _init_sharded would set - pssh.host_list = self.host_list - pssh.reachable_hosts = self.host_list - pssh.sharder = mock_sharder - pssh._use_process_sharding = True # Enable sharding for this test - # Add attributes needed by _shard_init_kwargs() - pssh.user = "test" - pssh.password = None - pssh.pkey = "id_rsa" - pssh.host_key_check = False - pssh.stop_on_errors = True - pssh.env_vars = None - - result = pssh.download_file("/remote/test.txt", "/tmp/test.txt", suffix_separator="_") - - # Verify sharder methods were called - mock_sharder.chunk_hosts.assert_called_once_with(self.host_list) - mock_sharder.execute_sharded.assert_called_once() - # download_file should return merged host->path dict - expected_result = {"host1": "/tmp/test.txt_host1", "host2": "/tmp/test.txt_host2"} - self.assertEqual(result, expected_result) + # Mock execute_sharded to return realistic download results + self.mock_execute_sharded.return_value = [ + {'result': {"host1": "/tmp/test.txt_host1"}, 'reachable_hosts': ["host1"], 'unreachable_hosts': []}, + {'result': {"host2": "/tmp/test.txt_host2"}, 'reachable_hosts': ["host2"], 'unreachable_hosts': []}, + ] + + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.download_file("/remote/test.txt", "/tmp/test.txt", suffix_separator="_") + + # Verify execute_sharded was called with real payloads + self.mock_execute_sharded.assert_called_once() + payloads = self.mock_execute_sharded.call_args[0][0] + + # Check that real payload was created correctly + self.assertEqual(len(payloads), 2) # Two shards + first_payload = payloads[0] + self.assertEqual(first_payload['operation'], 'download_file') + self.assertEqual(first_payload['remote_file'], '/remote/test.txt') + self.assertEqual(first_payload['local_file'], '/tmp/test.txt') + self.assertEqual(first_payload['suffix_separator'], '_') + self.assertIn('init', first_payload) + + # download_file should return merged host->path dict + expected_result = {"host1": "/tmp/test.txt_host1", "host2": "/tmp/test.txt_host2"} + self.assertEqual(result, expected_result) diff --git a/cvs/lib/parallel/unittests/test_pssh_sharder.py b/cvs/lib/parallel/unittests/test_pssh_sharder.py index 13f37752..ced6b6f3 100644 --- a/cvs/lib/parallel/unittests/test_pssh_sharder.py +++ b/cvs/lib/parallel/unittests/test_pssh_sharder.py @@ -230,7 +230,7 @@ def test_pssh_shard_worker_exec_mode(self, mock_pssh_class): mock_pssh_class.assert_called_once_with( log=None, host_list=['host1', 'host2'], user='test', process_output=False ) - mock_shard.exec.assert_called_once_with('echo hello', timeout=30, print_console=False) + mock_shard.exec.assert_called_once_with(cmd='echo hello', timeout=30, print_console=False) mock_shard.destroy_clients.assert_called_once() @patch('cvs.lib.parallel.pssh_sharder.Pssh') @@ -255,19 +255,19 @@ def test_pssh_shard_worker_cmd_list_mode(self, mock_pssh_class): # Verify direct operation call mock_pssh_class.assert_called_once_with(log=None, host_list=['host1'], user='test', process_output=False) - mock_shard.exec_cmd_list.assert_called_once_with(['echo 1', 'echo 2'], timeout=None, print_console=False) + mock_shard.exec_cmd_list.assert_called_once_with(cmd_list=['echo 1', 'echo 2'], print_console=False) mock_shard.destroy_clients.assert_called_once() @patch('cvs.lib.parallel.pssh_sharder.Pssh') - def test_pssh_shard_worker_scp_mode(self, mock_pssh_class): - """Test shard worker in scp mode with direct operation calls.""" + def test_pssh_shard_worker_upload_file_mode(self, mock_pssh_class): + """Test shard worker in upload_file mode with direct operation calls.""" mock_shard = MagicMock() mock_pssh_class.return_value = mock_shard mock_shard.reachable_hosts = ['host1'] mock_shard.unreachable_hosts = [] payload = { - 'operation': 'scp_file', + 'operation': 'upload_file', 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, 'local_file': 'test.txt', 'remote_file': '/tmp/test.txt', @@ -280,7 +280,36 @@ def test_pssh_shard_worker_scp_mode(self, mock_pssh_class): self.assertEqual(result, expected) # Verify direct operation call - mock_shard.scp_file.assert_called_once_with('test.txt', '/tmp/test.txt', recurse=False) + mock_shard.upload_file.assert_called_once_with( + local_file='test.txt', remote_file='/tmp/test.txt', recurse=False + ) + + @patch('cvs.lib.parallel.pssh_sharder.Pssh') + def test_pssh_shard_worker_download_file_mode(self, mock_pssh_class): + """Test shard worker in download_file mode with direct operation calls.""" + mock_shard = MagicMock() + mock_pssh_class.return_value = mock_shard + mock_shard.reachable_hosts = ['host1'] + mock_shard.unreachable_hosts = [] + mock_shard.download_file.return_value = {'host1': '/local/test_host1.txt'} + + payload = { + 'operation': 'download_file', + 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, + 'remote_file': '/tmp/test.txt', + 'local_file': 'test.txt', + 'recurse': False, + } + + result = PsshSharder.run_shard(payload) + + expected = {'result': {'host1': '/local/test_host1.txt'}, 'reachable_hosts': ['host1'], 'unreachable_hosts': []} + self.assertEqual(result, expected) + + # Verify direct operation call + mock_shard.download_file.assert_called_once_with( + remote_file='/tmp/test.txt', local_file='test.txt', recurse=False + ) @patch('cvs.lib.parallel.pssh_sharder.Pssh') def test_pssh_shard_worker_reboot_mode(self, mock_pssh_class): @@ -320,292 +349,23 @@ def test_pssh_shard_worker_unknown_operation(self, mock_pssh_class): self.assertIn('Unknown operation: unknown_operation', str(cm.exception)) @patch('cvs.lib.parallel.pssh_sharder.Pssh') - def test_pssh_shard_worker_registry_handler_mismatch(self, mock_pssh_class): - """Test shard worker with operation in registry but missing handler (implementation error).""" + def test_pssh_shard_worker_abc_enforcement(self, mock_pssh_class): + """Test shard worker properly enforces ABC operations.""" mock_shard = MagicMock() mock_pssh_class.return_value = mock_shard - # Temporarily add an operation to registry without corresponding handler - from cvs.lib.parallel.pssh_sharder import SUPPORTED_OPERATIONS - - original_operations = SUPPORTED_OPERATIONS.copy() - - try: - # Add a fake operation to registry - SUPPORTED_OPERATIONS['fake_operation'] = { - 'required_params': [], - 'optional_params': [], - 'returns_result': False, - } - - payload = { - 'operation': 'fake_operation', - 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, - } - - with self.assertRaises(RuntimeError) as cm: - PsshSharder.run_shard(payload) - - # Verify it's identified as a implementation error - self.assertIn('Implementation bug', str(cm.exception)) - self.assertIn('fake_operation', str(cm.exception)) - self.assertIn('supported (in registry) but handler is missing', str(cm.exception)) - - finally: - # Restore original registry - SUPPORTED_OPERATIONS.clear() - SUPPORTED_OPERATIONS.update(original_operations) - - -class TestPsshSharderApiCompliance(unittest.TestCase): - """ - Test suite to ensure PsshSharder supports all operations and parameters - required by MultiProcessPssh, preventing API drift issues. - """ - - def test_sharder_supports_all_multiprocess_operations(self): - """ - Verify that pssh_sharder supports all operations used by MultiProcessPssh. - - Uses the operation registry as the single source of truth instead of - fragile regex parsing. - """ - import re - import inspect - from cvs.lib.parallel.pssh_sharder import SUPPORTED_OPERATIONS - - # Get MultiProcessPssh source code to extract sharder operations - from cvs.lib.parallel.multiprocess_pssh import MultiProcessPssh - - multiprocess_source = inspect.getsource(MultiProcessPssh) - - # Extract operation names from create_payloads calls - operation_pattern = r"create_payloads\s*\(\s*['\"](\w+)['\"]" - used_operations = set(re.findall(operation_pattern, multiprocess_source)) - - # Get supported operations from registry (single source of truth) - supported_operations = set(SUPPORTED_OPERATIONS.keys()) - - # Find missing operations - missing_operations = used_operations - supported_operations - - if missing_operations: - error_msg = ( - f"PsshSharder missing support for {len(missing_operations)} operation(s) " - f"used by MultiProcessPssh:\n" - f"Missing operations: {sorted(missing_operations)}\n" - f"Used by MultiProcessPssh: {sorted(used_operations)}\n" - f"Supported by PsshSharder: {sorted(supported_operations)}\n\n" - f"Add missing operations to SUPPORTED_OPERATIONS registry in pssh_sharder.py" - ) - self.fail(error_msg) - - print(f"✓ Sharder Compatibility Check Passed: All {len(used_operations)} operations supported") - - def test_sharder_operations_can_execute_without_errors(self): - """ - Test that all sharder operations can execute without parameter/method errors. - - Uses the operation registry to dynamically generate test payloads, - ensuring all operations are tested consistently. - """ - from cvs.lib.parallel.pssh_sharder import SUPPORTED_OPERATIONS - - config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) - sharder = PsshSharder(config) - - # Dynamically generate test payloads from operation registry - test_payloads = {} - - for operation_name, op_config in SUPPORTED_OPERATIONS.items(): - payload = { - 'operation': operation_name, - 'init': {'log': None, 'host_list': ['testhost'], 'user': 'test'}, - } + # Test with an operation not in ABC + payload = { + 'operation': 'fake_operation_not_in_abc', + 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, + } - # Add required parameters with test values - param_values = { - 'cmd': 'echo test', - 'cmd_list': ['echo test'], - 'timeout': 30, - 'detailed': False, - 'local_file': '/tmp/test.txt', - 'remote_file': '/remote/test.txt', - 'recurse': False, - 'suffix_separator': '_', - } + with self.assertRaises(ValueError) as cm: + PsshSharder.run_shard(payload) - for param in op_config['required_params']: - if param in param_values: - payload[param] = param_values[param] - - # Add some optional parameters for testing - for param in op_config['optional_params']: - if param in param_values: - payload[param] = param_values[param] - - test_payloads[operation_name] = payload - - # Mock all Pssh methods to avoid actual SSH calls - with patch('cvs.lib.parallel.pssh_sharder.Pssh') as mock_pssh_class: - mock_shard = MagicMock() - mock_pssh_class.return_value = mock_shard - - # Configure return values - mock_shard.exec.return_value = {'testhost': 'output'} - mock_shard.exec_cmd_list.return_value = {'testhost': 'output'} - mock_shard.download_file.return_value = {'testhost': '/tmp/test.txt_testhost'} - mock_shard.reachable_hosts = ['testhost'] - mock_shard.unreachable_hosts = [] - - failed_operations = [] - - for operation_name, payload in test_payloads.items(): - try: - result = sharder.run_shard(payload) - - # Verify result structure - self.assertIn('result', result) - self.assertIn('reachable_hosts', result) - self.assertIn('unreachable_hosts', result) - - # Verify correct method was called on mock shard - if operation_name == 'exec': - mock_shard.exec.assert_called() - elif operation_name == 'cmd_list': - mock_shard.exec_cmd_list.assert_called() - elif operation_name == 'scp': - mock_shard.scp_file.assert_called() - elif operation_name == 'upload': - mock_shard.upload_file.assert_called() - elif operation_name == 'download': - mock_shard.download_file.assert_called() - elif operation_name == 'reboot': - mock_shard.reboot_connections.assert_called() - - except Exception as e: - failed_operations.append( - {'operation': operation_name, 'error': str(e), 'error_type': type(e).__name__} - ) - finally: - # Reset mocks for next iteration - mock_shard.reset_mock() - - if failed_operations: - error_msg = "Sharder operations failed with runtime errors:\n" - for failure in failed_operations: - error_msg += f" {failure['operation']}: {failure['error_type']}: {failure['error']}\n" - error_msg += "\nThis suggests missing parameters, wrong method names, or incorrect implementation." - self.fail(error_msg) - - print(f"✓ Runtime Compatibility Check Passed: All {len(test_payloads)} operations execute correctly") - - def test_operation_registry_completeness(self): - """ - Test that the operation registry is complete and consistent. - - Validates that all operations have proper configuration and that - the referenced Pssh methods actually exist. - """ - from cvs.lib.parallel.pssh_sharder import SUPPORTED_OPERATIONS - from cvs.lib.parallel.pssh import Pssh - import inspect - - errors = [] - - for operation_name, op_config in SUPPORTED_OPERATIONS.items(): - # Check required fields - required_fields = ['required_params', 'optional_params', 'returns_result'] - for field in required_fields: - if field not in op_config: - errors.append(f"Operation '{operation_name}' missing required field '{field}'") - - # Check that operation name corresponds to an actual method in Pssh class - if not hasattr(Pssh, operation_name): - errors.append(f"Operation '{operation_name}' does not correspond to a method in Pssh class") - - # Check method signature compatibility - if hasattr(Pssh, operation_name): - try: - method = getattr(Pssh, operation_name) - sig = inspect.signature(method) - method_params = set(sig.parameters.keys()) - {'self'} # Exclude 'self' - - config_params = set(op_config['required_params'] + op_config['optional_params']) - - # Check if all config parameters exist in method signature - missing_in_method = config_params - method_params - if missing_in_method: - errors.append( - f"Operation '{operation_name}': method missing parameters {sorted(missing_in_method)}" - ) - - except Exception as e: - errors.append(f"Operation '{operation_name}': Could not validate method signature: {e}") - - if errors: - error_msg = "Operation registry validation failed:\n" - for error in errors: - error_msg += f" - {error}\n" - error_msg += "\nFix the SUPPORTED_OPERATIONS registry in pssh_sharder.py" - self.fail(error_msg) - - print(f"✓ Operation Registry Validation Passed: All {len(SUPPORTED_OPERATIONS)} operations properly configured") - - def test_registry_handler_consistency(self): - """ - Test that all operations in registry have corresponding handlers in run_shard. - - This prevents the implementation error scenario where operations are added to - the registry but corresponding handler code is forgotten. - """ - from cvs.lib.parallel.pssh_sharder import SUPPORTED_OPERATIONS - import inspect - import re - - # Get the run_shard method source to extract handled operations - from cvs.lib.parallel.pssh_sharder import PsshSharder - - run_shard_source = inspect.getsource(PsshSharder.run_shard) - - # Extract operation names from if/elif operation == 'name' and operation in ['name1', 'name2'] patterns - # Pattern 1: operation == 'name' - equals_pattern = r"operation == ['\"](\w+)['\"]" - equals_operations = set(re.findall(equals_pattern, run_shard_source)) - - # Pattern 2: operation in ['name1', 'name2', ...] - in_pattern = r"operation in \[(.*?)\]" - in_matches = re.findall(in_pattern, run_shard_source) - in_operations = set() - for match in in_matches: - # Extract individual operation names from the list - names = re.findall(r"['\"](\w+)['\"]", match) - in_operations.update(names) - - handled_operations = equals_operations | in_operations - - # Operations in registry but not in handlers (implementation error) - registry_operations = set(SUPPORTED_OPERATIONS.keys()) - missing_handlers = registry_operations - handled_operations - - # Operations in handlers but not in registry (less critical, but good to know) - extra_handlers = handled_operations - registry_operations - - errors = [] - if missing_handlers: - errors.append(f"Operations in registry but missing handlers: {sorted(missing_handlers)}") - - if extra_handlers: - errors.append(f"Operations with handlers but not in registry: {sorted(extra_handlers)}") - - if errors: - error_msg = "Registry-Handler consistency check failed:\n" - for error in errors: - error_msg += f" - {error}\n" - error_msg += "\nEnsure registry and handler code are synchronized." - self.fail(error_msg) - - print(f"✓ Registry-Handler Consistency Check Passed: All {len(registry_operations)} operations have handlers") + # Verify proper error message + self.assertIn('Unknown operation: fake_operation_not_in_abc', str(cm.exception)) + self.assertIn('Supported:', str(cm.exception)) if __name__ == "__main__":