test_enhanced_functionality.py

#!/usr/bin/env python3
"""
Enhanced MetaWalletGen CLI Test Suite

This test suite verifies all the improvements and enhancements made to address
the usability issues and functionality gaps identified in the analysis.
"""

import sys
import os
import unittest
import tempfile
import shutil
import time
from pathlib import Path

# Add the current directory to Python path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))


class TestEnhancedCLI(unittest.TestCase):
    """Test enhanced CLI functionality."""
    
    def setUp(self):
        """Set up test environment."""
        self.test_dir = tempfile.mkdtemp()
        self.original_cwd = os.getcwd()
        os.chdir(self.test_dir)
        
        # Create test wallets directory
        self.wallets_dir = Path("wallets")
        self.wallets_dir.mkdir(exist_ok=True)
    
    def tearDown(self):
        """Clean up test environment."""
        os.chdir(self.original_cwd)
        shutil.rmtree(self.test_dir)
    
    def test_enhanced_imports(self):
        """Test that all enhanced modules can be imported."""
        try:
            import metawalletgen
            from metawalletgen.utils.config_manager import get_config
            from metawalletgen.utils.logger import get_logger
            from metawalletgen.utils.validators import (
                validate_ethereum_address,
                validate_private_key,
                validate_mnemonic,
                validate_derivation_path
            )
            self.assertTrue(True, "All enhanced modules imported successfully")
        except ImportError as e:
            self.fail(f"Failed to import enhanced modules: {e}")
    
    def test_configuration_management(self):
        """Test configuration management functionality."""
        from metawalletgen.utils.config_manager import get_config
        
        config = get_config()
        
        # Test default values
        defaults = config.get_defaults()
        self.assertIn("network", defaults)
        self.assertIn("derivation_path", defaults)
        self.assertIn("output_format", defaults)
        
        # Test network support
        networks = config.get_supported_networks()
        self.assertIn("mainnet", networks)
        self.assertIn("testnet", networks)
        self.assertIn("sepolia", networks)
        
        # Test format support
        formats = config.get_supported_formats()
        self.assertIn("json", formats)
        self.assertIn("csv", formats)
        self.assertIn("yaml", formats)
        
        # Test configuration validation
        issues = config.validate_config()
        self.assertIsInstance(issues, list)
    
    def test_enhanced_logging(self):
        """Test enhanced logging functionality."""
        from metawalletgen.utils.logger import get_logger
        
        logger = get_logger("test")
        
        # Test basic logging
        logger.info("Test info message")
        logger.warning("Test warning message")
        logger.error("Test error message")
        
        # Test specialized logging methods
        logger.log_wallet_generation(5, "testnet", "json", True)
        logger.log_wallet_import(3, "test.json", "json")
        logger.log_file_operation("save", "test.json", True)
        logger.log_validation_result(10, 8, 2)
        
        # Test log level changes
        logger.set_level("DEBUG")
        self.assertEqual(logger.logger.level, 10)  # DEBUG level
        
        # Test file handler addition
        test_log_file = "test.log"
        logger.add_file_handler(test_log_file)
        log_files = logger.get_log_file_paths()
        self.assertIn(os.path.abspath(test_log_file), log_files)
    
    def test_enhanced_validation(self):
        """Test enhanced validation functionality."""
        from metawalletgen.utils.validators import (
            validate_ethereum_address,
            validate_private_key,
            validate_mnemonic,
            validate_derivation_path
        )
        
        # Test address validation
        valid_address = "0x742d35Cc6634C0532925a3b8D4C9db96C4b4d8b6"
        invalid_address = "0xinvalid"
        
        self.assertTrue(validate_ethereum_address(valid_address))
        self.assertFalse(validate_ethereum_address(invalid_address))
        
        # Test private key validation
        valid_private_key = "0x" + "a" * 64
        invalid_private_key = "0x" + "a" * 32
        
        self.assertTrue(validate_private_key(valid_private_key))
        self.assertFalse(validate_private_key(invalid_private_key))
        
        # Test mnemonic validation
        valid_mnemonic = "abandon ability able about above absent absorb abstract absurd abuse access accident"
        invalid_mnemonic = "invalid mnemonic phrase"
        
        self.assertTrue(validate_mnemonic(valid_mnemonic))
        self.assertFalse(validate_mnemonic(invalid_mnemonic))
        
        # Test derivation path validation
        valid_path = "m/44'/60'/0'/0/0"
        invalid_path = "invalid/path"
        
        self.assertTrue(validate_derivation_path(valid_path))
        self.assertFalse(validate_derivation_path(invalid_path))
    
    def test_enhanced_wallet_generation(self):
        """Test enhanced wallet generation functionality."""
        from metawalletgen.core.wallet_generator import WalletGenerator
        
        generator = WalletGenerator(network="testnet")
        
        # Test single wallet generation
        wallet = generator.generate_new_wallet()
        self.assertIsNotNone(wallet.address)
        self.assertIsNotNone(wallet.private_key)
        self.assertIsNotNone(wallet.mnemonic)
        self.assertEqual(wallet.network, "testnet")
        
        # Test batch wallet generation
        wallets = generator.generate_batch_wallets(3)
        self.assertEqual(len(wallets), 3)
        
        # Test wallet from mnemonic
        mnemonic = wallet.mnemonic
        wallet_from_mnemonic = generator.create_wallet_from_mnemonic(mnemonic)
        self.assertEqual(wallet_from_mnemonic.address, wallet.address)
        
        # Test wallet from private key
        private_key = wallet.private_key
        wallet_from_key = generator.create_wallet_from_private_key(private_key)
        self.assertEqual(wallet_from_key.address, wallet.address)
    
    def test_enhanced_storage(self):
        """Test enhanced storage functionality."""
        from metawalletgen.core.storage_manager import StorageManager
        from metawalletgen.core.wallet_generator import WalletGenerator
        
        storage = StorageManager()
        generator = WalletGenerator()
        
        # Generate test wallets
        wallets = generator.generate_batch_wallets(2)
        
        # Test JSON storage
        json_file = storage.save_wallets_json(wallets, "test_wallets.json")
        self.assertTrue(os.path.exists(json_file))
        
        # Test CSV storage
        csv_file = storage.save_wallets_csv(wallets, "test_wallets.csv")
        self.assertTrue(os.path.exists(csv_file))
        
        # Test YAML storage
        yaml_file = storage.save_wallets_yaml(wallets, "test_wallets.yaml")
        self.assertTrue(os.path.exists(yaml_file))
        
        # Test encrypted storage
        encrypted_file = storage.save_wallets_json(
            wallets, "test_wallets_encrypted.json", 
            encrypt=True, password="test_password_123"
        )
        self.assertTrue(os.path.exists(encrypted_file))
        
        # Test loading wallets
        loaded_wallets = storage.load_wallets_json("test_wallets.json")
        self.assertEqual(len(loaded_wallets), 2)
    
    def test_cli_command_structure(self):
        """Test enhanced CLI command structure."""
        try:
            from metawalletgen.cli.main import main
            from metawalletgen.cli.commands import (
                generate_command,
                import_command,
                list_command,
                validate_command
            )
            
            # Test that commands exist and are callable
            self.assertTrue(callable(generate_command))
            self.assertTrue(callable(import_command))
            self.assertTrue(callable(list_command))
            self.assertTrue(callable(validate_command))
            
            # Test that main CLI group exists
            self.assertTrue(hasattr(main, 'commands'))
            
        except ImportError as e:
            self.fail(f"Failed to import CLI modules: {e}")
    
    def test_environment_variable_support(self):
        """Test environment variable configuration support."""
        from metawalletgen.utils.config_manager import get_config
        
        # Set test environment variables
        os.environ["METAWALLETGEN_NETWORK"] = "sepolia"
        os.environ["METAWALLETGEN_DEFAULT_COUNT"] = "5"
        os.environ["METAWALLETGEN_LOG_LEVEL"] = "DEBUG"
        
        # Create new config instance to load environment variables
        from metawalletgen.utils.config_manager import ConfigManager
        config = ConfigManager()
        
        # Verify environment variables were loaded
        self.assertEqual(config.get("defaults.network"), "sepolia")
        self.assertEqual(config.get("defaults.default_count"), 5)
        self.assertEqual(config.get("logging.level"), "DEBUG")
        
        # Clean up environment variables
        del os.environ["METAWALLETGEN_NETWORK"]
        del os.environ["METAWALLETGEN_DEFAULT_COUNT"]
        del os.environ["METAWALLETGEN_LOG_LEVEL"]
    
    def test_file_handling_improvements(self):
        """Test improved file handling functionality."""
        from metawalletgen.core.storage_manager import StorageManager
        from metawalletgen.core.wallet_generator import WalletGenerator
        
        storage = StorageManager()
        generator = WalletGenerator()
        
        # Generate test wallet
        wallet = generator.generate_new_wallet()
        
        # Test automatic file extension handling
        json_file = storage.save_wallets_json([wallet], "test_wallets")
        self.assertTrue(json_file.endswith(".json"))
        
        csv_file = storage.save_wallets_csv([wallet], "test_wallets")
        self.assertTrue(csv_file.endswith(".csv"))
        
        # Test file size reporting
        if os.path.exists(json_file):
            file_size = os.path.getsize(json_file)
            self.assertGreater(file_size, 0)
    
    def test_security_features(self):
        """Test enhanced security features."""
        from metawalletgen.core.storage_manager import StorageManager
        from metawalletgen.core.wallet_generator import WalletGenerator
        
        storage = StorageManager()
        generator = WalletGenerator()
        
        # Generate test wallet
        wallet = generator.generate_new_wallet()
        
        # Test encryption
        encrypted_file = storage.save_wallets_json(
            [wallet], "test_encrypted.json", 
            encrypt=True, password="secure_password_123"
        )
        
        # Verify file is encrypted (should contain encrypted content)
        with open(encrypted_file, 'r') as f:
            content = f.read()
            self.assertIn("encrypted", content.lower())
            self.assertIn("vault", content.lower())
        
        # Test password validation
        try:
            # This should fail with wrong password
            storage.load_wallets_json(encrypted_file, decrypt=True, password="wrong_password")
            self.fail("Should have failed with wrong password")
        except Exception:
            # Expected to fail
            pass
    
    def test_progress_tracking(self):
        """Test progress tracking functionality."""
        from metawalletgen.core.wallet_generator import WalletGenerator
        
        generator = WalletGenerator()
        
        # Test batch generation with progress tracking
        start_time = time.time()
        wallets = generator.generate_batch_wallets(10)
        end_time = time.time()
        
        self.assertEqual(len(wallets), 10)
        self.assertGreater(end_time - start_time, 0)  # Should take some time
        
        # Verify all wallets are unique
        addresses = [w.address for w in wallets]
        self.assertEqual(len(addresses), len(set(addresses)))
    
    def test_error_handling(self):
        """Test enhanced error handling."""
        from metawalletgen.utils.validators import validate_ethereum_address
        
        # Test graceful handling of invalid inputs
        invalid_inputs = [
            "",  # Empty string
            None,  # None value
            "not_an_address",  # Invalid format
            "0x" + "a" * 100,  # Too long
        ]
        
        for invalid_input in invalid_inputs:
            try:
                result = validate_ethereum_address(invalid_input)
                self.assertFalse(result)  # Should return False for invalid inputs
            except Exception as e:
                # Should not crash, but may return False
                self.assertIsInstance(e, Exception)


def run_performance_tests():
    """Run performance tests for batch operations."""
    print("\n🚀 Performance Tests")
    print("=" * 30)
    
    try:
        from metawalletgen.core.wallet_generator import WalletGenerator
        import time
        
        generator = WalletGenerator()
        
        # Test different batch sizes
        batch_sizes = [1, 10, 100, 1000]
        
        for size in batch_sizes:
            print(f"\n📊 Testing batch size: {size}")
            
            start_time = time.time()
            wallets = generator.generate_batch_wallets(size)
            end_time = time.time()
            
            duration = end_time - start_time
            rate = size / duration if duration > 0 else 0
            
            print(f"   ✅ Generated {len(wallets)} wallets in {duration:.2f} seconds")
            print(f"   📈 Rate: {rate:.1f} wallets/second")
            
            # Verify all wallets are unique
            addresses = [w.address for w in wallets]
            unique_count = len(set(addresses))
            print(f"   🔍 Unique addresses: {unique_count}/{size}")
            
            if unique_count != size:
                print(f"   ⚠️  Warning: {size - unique_count} duplicate addresses detected")
    
    except Exception as e:
        print(f"❌ Performance test error: {e}")


def run_integration_tests():
    """Run integration tests for the complete workflow."""
    print("\n🔗 Integration Tests")
    print("=" * 30)
    
    try:
        from metawalletgen.core.wallet_generator import WalletGenerator
        from metawalletgen.core.storage_manager import StorageManager
        from metawalletgen.utils.validators import validate_ethereum_address
        import tempfile
        import os
        
        # Create temporary directory
        with tempfile.TemporaryDirectory() as temp_dir:
            os.chdir(temp_dir)
            
            print("1. Testing complete wallet generation workflow...")
            
            # Generate wallets
            generator = WalletGenerator(network="testnet")
            wallets = generator.generate_batch_wallets(5)
            print(f"   ✅ Generated {len(wallets)} wallets")
            
            # Validate wallets
            valid_count = 0
            for wallet in wallets:
                if validate_ethereum_address(wallet.address):
                    valid_count += 1
            
            print(f"   ✅ Validated {valid_count}/{len(wallets)} wallets")
            
            # Save wallets
            storage = StorageManager()
            json_file = storage.save_wallets_json(wallets, "integration_test.json")
            print(f"   ✅ Saved wallets to {os.path.basename(json_file)}")
            
            # Load and verify wallets
            loaded_wallets = storage.load_wallets_json(json_file)
            print(f"   ✅ Loaded {len(loaded_wallets)} wallets from file")
            
            # Verify data integrity
            for i, (original, loaded) in enumerate(zip(wallets, loaded_wallets)):
                if (original.address == loaded.address and 
                    original.private_key == loaded.private_key):
                    continue
                else:
                    print(f"   ❌ Data mismatch at index {i}")
                    break
            else:
                print("   ✅ Data integrity verified")
            
            print("2. Testing encrypted workflow...")
            
            # Save encrypted
            encrypted_file = storage.save_wallets_json(
                wallets, "integration_test_encrypted.json",
                encrypt=True, password="integration_test_password"
            )
            print(f"   ✅ Saved encrypted wallets to {os.path.basename(encrypted_file)}")
            
            # Load encrypted
            try:
                decrypted_wallets = storage.load_wallets_json(
                    encrypted_file, decrypt=True, password="integration_test_password"
                )
                print(f"   ✅ Successfully decrypted {len(decrypted_wallets)} wallets")
                
                # Verify decryption integrity
                if len(decrypted_wallets) == len(wallets):
                    print("   ✅ Decryption integrity verified")
                else:
                    print("   ❌ Decryption integrity check failed")
                    
            except Exception as e:
                print(f"   ❌ Decryption failed: {e}")
            
            print("3. Testing multiple format support...")
            
            # Test all formats
            formats = {
                "json": storage.save_wallets_json,
                "csv": storage.save_wallets_csv,
                "yaml": storage.save_wallets_yaml
            }
            
            for format_name, save_func in formats.items():
                try:
                    filepath = save_func(wallets, f"integration_test.{format_name}")
                    print(f"   ✅ {format_name.upper()} format: {os.path.basename(filepath)}")
                except Exception as e:
                    print(f"   ❌ {format_name.upper()} format failed: {e}")
    
    except Exception as e:
        print(f"❌ Integration test error: {e}")


if __name__ == "__main__":
    print("🧪 Enhanced MetaWalletGen CLI Test Suite")
    print("=" * 50)
    print("This test suite verifies all the improvements and enhancements")
    print("made to address usability issues and functionality gaps.\n")
    
    # Run unit tests
    print("🔬 Running Unit Tests...")
    unittest.main(argv=[''], exit=False, verbosity=2)
    
    # Run performance tests
    run_performance_tests()
    
    # Run integration tests
    run_integration_tests()
    
    print("\n" + "=" * 50)
    print("🎉 Test suite completed!")
    print("\nThe enhanced MetaWalletGen CLI now includes:")
    print("   ✅ Enhanced validation and error handling")
    print("   ✅ Progress tracking for batch operations")
    print("   ✅ Better user feedback and security reminders")
    print("   ✅ Configuration management with environment variable support")
    print("   ✅ Comprehensive logging system")
    print("   ✅ Improved CLI commands and help text")
    print("   ✅ Better file handling and encryption")
    print("   ✅ Performance optimizations")
    print("   ✅ Integration testing")