CIA/e-voting-system/tests/test_poa_implementation.py

"""
Comprehensive Test Suite for PoA Blockchain Implementation

Tests the bootnode and validator services without requiring Docker.
Verifies core logic, consensus, and data structures.
"""

import sys
import json
import hashlib
import pytest
from pathlib import Path
from unittest.mock import Mock, AsyncMock, patch, MagicMock

# Add paths for imports
sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent.parent / "bootnode"))
sys.path.insert(0, str(Path(__file__).parent.parent / "validator"))

# ============================================================================
# BOOTNODE TESTS
# ============================================================================

class TestBootnodeService:
    """Test Bootnode peer discovery service"""

    def test_bootnode_initialization(self):
        """Test bootnode peer registry initialization"""
        from bootnode.bootnode import PeerRegistry

        registry = PeerRegistry(peer_timeout_seconds=300)

        assert registry.peers == {}
        assert registry.peer_timeout == 300
        print("✅ Bootnode initialization test passed")

    def test_peer_registration(self):
        """Test registering a peer with bootnode"""
        from bootnode.bootnode import PeerRegistry, PeerInfo

        registry = PeerRegistry()

        peer = PeerInfo(
            node_id="validator-1",
            ip="validator-1",
            p2p_port=30303,
            rpc_port=8001,
            public_key="0x1234567890abcdef"
        )

        registry.register_peer(peer)

        assert "validator-1" in registry.peers
        assert registry.peers["validator-1"]["info"].node_id == "validator-1"
        print("✅ Peer registration test passed")

    def test_peer_discovery(self):
        """Test discovering peers from bootnode"""
        from bootnode.bootnode import PeerRegistry, PeerInfo

        registry = PeerRegistry()

        # Register 3 peers
        for i in range(1, 4):
            peer = PeerInfo(
                node_id=f"validator-{i}",
                ip=f"validator-{i}",
                p2p_port=30300 + i,
                rpc_port=8000 + i,
                public_key=f"0x{'0' * 40}"
            )
            registry.register_peer(peer)

        # Discover peers (exclude validator-1)
        discovered = registry.get_peers_except("validator-1")

        assert len(discovered) == 2
        node_ids = [p.node_id for p in discovered]
        assert "validator-2" in node_ids
        assert "validator-3" in node_ids
        assert "validator-1" not in node_ids
        print("✅ Peer discovery test passed")

    def test_peer_heartbeat(self):
        """Test updating peer heartbeat"""
        from bootnode.bootnode import PeerRegistry, PeerInfo
        import time

        registry = PeerRegistry(peer_timeout_seconds=5)

        peer = PeerInfo(
            node_id="validator-1",
            ip="validator-1",
            p2p_port=30303,
            rpc_port=8001
        )

        registry.register_peer(peer)
        old_heartbeat = registry.peers["validator-1"]["last_heartbeat"]

        # Update heartbeat
        time.sleep(0.1)
        registry.update_heartbeat("validator-1")
        new_heartbeat = registry.peers["validator-1"]["last_heartbeat"]

        assert new_heartbeat > old_heartbeat
        print("✅ Peer heartbeat test passed")

    def test_stale_peer_cleanup(self):
        """Test cleanup of stale peers"""
        from bootnode.bootnode import PeerRegistry, PeerInfo
        import time

        registry = PeerRegistry(peer_timeout_seconds=1)

        # Register peer
        peer = PeerInfo(
            node_id="validator-1",
            ip="validator-1",
            p2p_port=30303,
            rpc_port=8001
        )
        registry.register_peer(peer)
        assert len(registry.get_all_peers()) == 1

        # Manually set old heartbeat time
        registry.peers["validator-1"]["last_heartbeat"] = time.time() - 2

        # Cleanup stale peers
        removed_count = registry.cleanup_stale_peers()

        assert removed_count == 1
        assert len(registry.get_all_peers()) == 0
        print("✅ Stale peer cleanup test passed")


# ============================================================================
# VALIDATOR/BLOCKCHAIN TESTS
# ============================================================================

class TestBlockchain:
    """Test Blockchain data structure and operations"""

    def test_genesis_block_creation(self):
        """Test genesis block is created correctly"""
        from validator.validator import Blockchain

        blockchain = Blockchain()

        assert blockchain.get_chain_length() == 1
        genesis = blockchain.get_block(0)
        assert genesis.index == 0
        assert genesis.validator == "genesis"
        assert genesis.transactions == []
        print("✅ Genesis block creation test passed")

    def test_block_hash_calculation(self):
        """Test block hash is calculated correctly"""
        from validator.validator import Blockchain, Block, Transaction

        blockchain = Blockchain()

        # Create a block
        tx = Transaction(
            voter_id="test-voter",
            election_id=1,
            encrypted_vote="0x1234",
            ballot_hash="0xabcd",
            timestamp=1699360000
        )

        block = Block(
            index=1,
            prev_hash=blockchain.get_latest_block().block_hash,
            timestamp=1699360010,
            transactions=[tx],
            validator="validator-1"
        )

        # Calculate hash
        hash1 = Blockchain.calculate_block_hash(block)

        # Same block should produce same hash
        hash2 = Blockchain.calculate_block_hash(block)

        assert hash1 == hash2
        assert hash1.startswith("0x")
        assert len(hash1) == 66  # "0x" + 64 hex chars
        print("✅ Block hash calculation test passed")

    def test_block_validation(self):
        """Test block validation logic"""
        from validator.validator import Blockchain, Block, Transaction

        blockchain = Blockchain()

        tx = Transaction(
            voter_id="test-voter",
            election_id=1,
            encrypted_vote="0x1234",
            ballot_hash="0xabcd",
            timestamp=1699360000
        )

        # Create valid block
        prev_block = blockchain.get_latest_block()
        block = Block(
            index=1,
            prev_hash=prev_block.block_hash,
            timestamp=1699360010,
            transactions=[tx],
            validator="validator-1"
        )
        block.block_hash = Blockchain.calculate_block_hash(block)

        assert blockchain.validate_block(block) == True
        print("✅ Valid block validation test passed")

        # Test invalid block (wrong index)
        invalid_block = Block(
            index=99,  # Wrong index
            prev_hash=prev_block.block_hash,
            timestamp=1699360010,
            transactions=[tx],
            validator="validator-1"
        )
        invalid_block.block_hash = Blockchain.calculate_block_hash(invalid_block)

        assert blockchain.validate_block(invalid_block) == False
        print("✅ Invalid block (wrong index) detection test passed")

        # Test invalid block (wrong prev_hash)
        invalid_block2 = Block(
            index=1,
            prev_hash="0x" + "0" * 64,  # Wrong prev_hash
            timestamp=1699360010,
            transactions=[tx],
            validator="validator-1"
        )
        invalid_block2.block_hash = Blockchain.calculate_block_hash(invalid_block2)

        assert blockchain.validate_block(invalid_block2) == False
        print("✅ Invalid block (wrong prev_hash) detection test passed")

        # Test invalid block (unauthorized validator)
        invalid_block3 = Block(
            index=1,
            prev_hash=prev_block.block_hash,
            timestamp=1699360010,
            transactions=[tx],
            validator="unauthorized-node"  # Not in authorized list
        )
        invalid_block3.block_hash = Blockchain.calculate_block_hash(invalid_block3)

        assert blockchain.validate_block(invalid_block3) == False
        print("✅ Invalid block (unauthorized validator) detection test passed")

    def test_add_block_to_chain(self):
        """Test adding valid blocks to blockchain"""
        from validator.validator import Blockchain, Block, Transaction

        blockchain = Blockchain()

        tx = Transaction(
            voter_id="test-voter",
            election_id=1,
            encrypted_vote="0x1234",
            ballot_hash="0xabcd",
            timestamp=1699360000
        )

        prev_block = blockchain.get_latest_block()
        block = Block(
            index=1,
            prev_hash=prev_block.block_hash,
            timestamp=1699360010,
            transactions=[tx],
            validator="validator-1"
        )
        block.block_hash = Blockchain.calculate_block_hash(block)

        result = blockchain.add_block(block)

        assert result == True
        assert blockchain.get_chain_length() == 2
        assert blockchain.get_block(1).index == 1
        print("✅ Add block to chain test passed")

    def test_blockchain_integrity_verification(self):
        """Test blockchain integrity verification"""
        from validator.validator import Blockchain, Block, Transaction

        blockchain = Blockchain()

        # Add valid blocks
        for i in range(1, 4):
            prev_block = blockchain.get_latest_block()
            tx = Transaction(
                voter_id=f"voter-{i}",
                election_id=1,
                encrypted_vote=f"0x{i:04x}",
                ballot_hash=f"0x{i:04x}",
                timestamp=1699360000 + i
            )

            block = Block(
                index=i,
                prev_hash=prev_block.block_hash,
                timestamp=1699360010 + i,
                transactions=[tx],
                validator=f"validator-{(i % 3) + 1}"
            )
            block.block_hash = Blockchain.calculate_block_hash(block)
            blockchain.add_block(block)

        # Verify integrity
        assert blockchain.verify_integrity() == True
        print("✅ Blockchain integrity verification test passed")

    def test_chain_immutability(self):
        """Test that modifying past blocks breaks chain"""
        from validator.validator import Blockchain, Block, Transaction

        blockchain = Blockchain()

        # Add a block
        prev_block = blockchain.get_latest_block()
        tx = Transaction(
            voter_id="voter-1",
            election_id=1,
            encrypted_vote="0x1234",
            ballot_hash="0xabcd",
            timestamp=1699360000
        )

        block = Block(
            index=1,
            prev_hash=prev_block.block_hash,
            timestamp=1699360010,
            transactions=[tx],
            validator="validator-1"
        )
        block.block_hash = Blockchain.calculate_block_hash(block)
        blockchain.add_block(block)

        # Verify chain is valid
        assert blockchain.verify_integrity() == True

        # Try to modify the transaction in block 1
        blockchain.chain[1].transactions[0].encrypted_vote = "0x9999"

        # Chain should now be invalid (hash mismatch)
        assert blockchain.verify_integrity() == False
        print("✅ Chain immutability test passed")


# ============================================================================
# POA CONSENSUS TESTS
# ============================================================================

class TestPoAConsensus:
    """Test Proof-of-Authority consensus mechanism"""

    def test_round_robin_block_creation(self):
        """Test round-robin block creation eligibility"""
        from validator.validator import PoAValidator

        # Create 3 validators
        validators = [
            PoAValidator("validator-1", "0x1234", "http://bootnode:8546", 8001, 30303),
            PoAValidator("validator-2", "0x5678", "http://bootnode:8546", 8002, 30304),
            PoAValidator("validator-3", "0xabcd", "http://bootnode:8546", 8003, 30305),
        ]

        # Test round-robin eligibility
        # Block 0: genesis (no validator)
        # Block 1: validator-1 eligible (1 % 3 == 0)
        # Block 2: validator-2 eligible (2 % 3 == 1)
        # Block 3: validator-3 eligible (3 % 3 == 2)
        # Block 4: validator-1 eligible (4 % 3 == 0)

        validator1 = validators[0]
        assert validator1.should_create_block() == True  # Block 1

        validator2 = validators[1]
        assert validator2.should_create_block() == False  # Not eligible for block 1

        validator3 = validators[2]
        assert validator3.should_create_block() == False  # Not eligible for block 1

        print("✅ Round-robin block creation test passed")

    def test_authorized_validators(self):
        """Test that only authorized validators can create blocks"""
        from validator.validator import Blockchain, Block, Transaction

        blockchain = Blockchain()

        # Test authorized validators
        authorized = Blockchain.AUTHORIZED_VALIDATORS
        assert "validator-1" in authorized
        assert "validator-2" in authorized
        assert "validator-3" in authorized

        # Create block with authorized validator
        tx = Transaction(
            voter_id="voter-1",
            election_id=1,
            encrypted_vote="0x1234",
            ballot_hash="0xabcd",
            timestamp=1699360000
        )

        prev_block = blockchain.get_latest_block()
        block = Block(
            index=1,
            prev_hash=prev_block.block_hash,
            timestamp=1699360010,
            transactions=[tx],
            validator="validator-1"
        )
        block.block_hash = Blockchain.calculate_block_hash(block)

        assert blockchain.validate_block(block) == True
        print("✅ Authorized validator test passed")

        # Create block with unauthorized validator
        block.validator = "unauthorized-node"
        block.block_hash = Blockchain.calculate_block_hash(block)

        assert blockchain.validate_block(block) == False
        print("✅ Unauthorized validator rejection test passed")


# ============================================================================
# DATA STRUCTURE TESTS
# ============================================================================

class TestDataStructures:
    """Test data models and serialization"""

    def test_transaction_model(self):
        """Test Transaction data model"""
        from validator.validator import Transaction

        tx = Transaction(
            voter_id="voter-123",
            election_id=1,
            encrypted_vote="0x1234567890abcdef",
            ballot_hash="0xabcdef1234567890",
            proof="0x...",
            timestamp=1699360000
        )

        assert tx.voter_id == "voter-123"
        assert tx.election_id == 1
        assert tx.encrypted_vote == "0x1234567890abcdef"
        print("✅ Transaction model test passed")

    def test_block_serialization(self):
        """Test block to/from dictionary conversion"""
        from validator.validator import Block, Transaction

        tx = Transaction(
            voter_id="voter-1",
            election_id=1,
            encrypted_vote="0x1234",
            ballot_hash="0xabcd",
            timestamp=1699360000
        )

        block = Block(
            index=1,
            prev_hash="0x" + "0" * 64,
            timestamp=1699360010,
            transactions=[tx],
            validator="validator-1",
            block_hash="0x" + "1" * 64,
            signature="0x" + "2" * 64
        )

        # Serialize to dict
        block_dict = block.to_dict()

        assert block_dict["index"] == 1
        assert block_dict["validator"] == "validator-1"
        assert len(block_dict["transactions"]) == 1
        assert block_dict["transactions"][0]["voter_id"] == "voter-1"

        # Deserialize from dict
        restored_block = Block.from_dict(block_dict)

        assert restored_block.index == block.index
        assert restored_block.validator == block.validator
        assert len(restored_block.transactions) == 1
        print("✅ Block serialization test passed")


# ============================================================================
# INTEGRATION TESTS
# ============================================================================

class TestIntegration:
    """Integration tests for the full system"""

    def test_multi_validator_consensus(self):
        """Test that multiple validators can reach consensus"""
        from validator.validator import Blockchain, Block, Transaction

        # Create 3 blockchains (one per validator)
        blockchains = [Blockchain(), Blockchain(), Blockchain()]

        # Simulate block creation and consensus
        for block_num in range(1, 4):
            # Determine which validator creates this block (round-robin)
            creator_idx = (block_num - 1) % 3
            validator_name = f"validator-{creator_idx + 1}"

            # Create block
            prev_block = blockchains[creator_idx].get_latest_block()
            tx = Transaction(
                voter_id=f"voter-{block_num}",
                election_id=1,
                encrypted_vote=f"0x{block_num:04x}",
                ballot_hash=f"0x{block_num:04x}",
                timestamp=1699360000 + block_num
            )

            block = Block(
                index=block_num,
                prev_hash=prev_block.block_hash,
                timestamp=1699360010 + block_num,
                transactions=[tx],
                validator=validator_name
            )
            block.block_hash = Blockchain.calculate_block_hash(block)

            # Add to creator's chain
            blockchains[creator_idx].add_block(block)

            # Broadcast to other validators
            for i in range(3):
                if i != creator_idx:
                    blockchains[i].add_block(block)

        # Verify all validators have same blockchain
        chain_hashes = []
        for i, bc in enumerate(blockchains):
            last_block_hash = bc.get_latest_block().block_hash
            chain_hashes.append(last_block_hash)
            assert bc.verify_integrity() == True

        # All should have same last block
        assert chain_hashes[0] == chain_hashes[1] == chain_hashes[2]
        print("✅ Multi-validator consensus test passed")

    def test_vote_immutability_across_validators(self):
        """Test that votes are immutable once recorded on blockchain"""
        from validator.validator import Blockchain, Block, Transaction

        blockchain = Blockchain()

        # Record a vote
        original_vote = Transaction(
            voter_id="voter-immutable",
            election_id=1,
            encrypted_vote="0x1234567890abcdef",
            ballot_hash="0xabcd",
            timestamp=1699360000
        )

        prev_block = blockchain.get_latest_block()
        block = Block(
            index=1,
            prev_hash=prev_block.block_hash,
            timestamp=1699360010,
            transactions=[original_vote],
            validator="validator-1"
        )
        block.block_hash = Blockchain.calculate_block_hash(block)
        blockchain.add_block(block)

        # Verify vote is recorded
        assert blockchain.chain[1].transactions[0].encrypted_vote == "0x1234567890abcdef"
        assert blockchain.verify_integrity() == True

        # Try to modify the vote (simulating attack)
        blockchain.chain[1].transactions[0].encrypted_vote = "0xhackedvalue"

        # Blockchain integrity should now fail
        assert blockchain.verify_integrity() == False
        print("✅ Vote immutability test passed")


# ============================================================================
# JSON-RPC TESTS
# ============================================================================

class TestJsonRPC:
    """Test JSON-RPC interface"""

    def test_json_rpc_structure(self):
        """Test JSON-RPC request/response format"""
        import json

        # Valid JSON-RPC request
        request = {
            "jsonrpc": "2.0",
            "method": "eth_sendTransaction",
            "params": [{"data": "0x1234"}],
            "id": 1
        }

        # Verify structure
        assert request["jsonrpc"] == "2.0"
        assert request["method"] in ["eth_sendTransaction", "eth_getTransactionReceipt", "eth_blockNumber", "eth_getBlockByNumber"]
        assert "id" in request

        # Valid JSON-RPC response
        response = {
            "jsonrpc": "2.0",
            "result": "0xabc123",
            "id": 1
        }

        assert response["jsonrpc"] == "2.0"
        assert "result" in response or "error" in response
        assert response["id"] == 1

        print("✅ JSON-RPC structure test passed")


# ============================================================================
# RUN TESTS
# ============================================================================

if __name__ == "__main__":
    print("\n" + "=" * 80)
    print("PoA BLOCKCHAIN IMPLEMENTATION TEST SUITE")
    print("=" * 80 + "\n")

    # Bootnode tests
    print("\n--- BOOTNODE TESTS ---\n")
    bootnode_tests = TestBootnodeService()
    bootnode_tests.test_bootnode_initialization()
    bootnode_tests.test_peer_registration()
    bootnode_tests.test_peer_discovery()
    bootnode_tests.test_peer_heartbeat()
    bootnode_tests.test_stale_peer_cleanup()

    # Blockchain tests
    print("\n--- BLOCKCHAIN TESTS ---\n")
    blockchain_tests = TestBlockchain()
    blockchain_tests.test_genesis_block_creation()
    blockchain_tests.test_block_hash_calculation()
    blockchain_tests.test_block_validation()
    blockchain_tests.test_add_block_to_chain()
    blockchain_tests.test_blockchain_integrity_verification()
    blockchain_tests.test_chain_immutability()

    # PoA Consensus tests
    print("\n--- POA CONSENSUS TESTS ---\n")
    consensus_tests = TestPoAConsensus()
    consensus_tests.test_round_robin_block_creation()
    consensus_tests.test_authorized_validators()

    # Data structure tests
    print("\n--- DATA STRUCTURE TESTS ---\n")
    data_tests = TestDataStructures()
    data_tests.test_transaction_model()
    data_tests.test_block_serialization()

    # Integration tests
    print("\n--- INTEGRATION TESTS ---\n")
    integration_tests = TestIntegration()
    integration_tests.test_multi_validator_consensus()
    integration_tests.test_vote_immutability_across_validators()

    # JSON-RPC tests
    print("\n--- JSON-RPC TESTS ---\n")
    jsonrpc_tests = TestJsonRPC()
    jsonrpc_tests.test_json_rpc_structure()

    print("\n" + "=" * 80)
    print("✅ ALL TESTS PASSED!")
    print("=" * 80 + "\n")

    print("Test Summary:")
    print("  ✅ Bootnode: 5/5 tests passed")
    print("  ✅ Blockchain: 6/6 tests passed")
    print("  ✅ PoA Consensus: 2/2 tests passed")
    print("  ✅ Data Structures: 2/2 tests passed")
    print("  ✅ Integration: 2/2 tests passed")
    print("  ✅ JSON-RPC: 1/1 tests passed")
    print("  ───────────────────────")
    print("  ✅ TOTAL: 18/18 tests passed\n")