test_chain_sync.py

import asyncio
import os
import time
import shutil
from src.core.chain import Blockchain
from src.network.node import Node

# Clean up previous tests
if os.path.exists("local_chain_8000.json"): os.remove("local_chain_8000.json")
if os.path.exists("local_chain_8001.json"): os.remove("local_chain_8001.json")

async def test_persistence():
    print("\n--- TEST 1: PERSISTENCE ---")
    chain = Blockchain()
    # We manually set the 'chain_file' conceptually by saving it manually or mocking Node
    # Let's use Node to drive it
    node = Node("127.0.0.1", 8000, chain)
    
    # Mine a block
    print("Mining block 1 on Node A...")
    chain.mine_pending_transactions("MinerA")
    node.blockchain.save_chain("local_chain_8000.json")
    
    # Verify file exists
    if os.path.exists("local_chain_8000.json"):
        print("Success: Chain file created.")
    else:
        print("Failure: Chain file not found.")
        return

    # Create new node/chain instance and load
    print("Restarting Node A (Loading chain)...")
    chain2 = Blockchain()
    chain2.load_chain("local_chain_8000.json")
    
    if len(chain2.chain) == 2: # Genesis + 1
        print("Success: Chain loaded with correct height.")
    else:
        print(f"Failure: Expected height 2, got {len(chain2.chain)}")

async def test_sync():
    print("\n--- TEST 2: SYNC ---")
    
    # Setup Node A (The Authority) - Already has 2 blocks from Test 1
    chainA = Blockchain()
    nodeA = Node("127.0.0.1", 8000, chainA) # Will load local_chain_8000.json automatically
    
    # Setup Node B (The Newcomer) - Empty (only Genesis)
    chainB = Blockchain()
    nodeB = Node("127.0.0.1", 8001, chainB)
    
    # Start Servers
    serverA = asyncio.create_task(nodeA.start_server())
    serverB = asyncio.create_task(nodeB.start_server())
    
    await asyncio.sleep(1)
    
    print(f"Node A Height: {len(nodeA.blockchain.chain)}") # Should be 2
    print(f"Node B Height: {len(nodeB.blockchain.chain)}") # Should be 1
    
    # Connect B to A
    print("Connecting Node B to Node A...")
    await nodeB.connect_to_peer("127.0.0.1", 8000)
    
    # Wait for sync (handshake -> query -> response -> process)
    await asyncio.sleep(2)
    
    print(f"Node B Height after sync: {len(nodeB.blockchain.chain)}")
    
    if len(nodeB.blockchain.chain) == len(nodeA.blockchain.chain):
        print("Success: Node B synced with Node A.")
    else:
        print("Failure: Node B did not sync.")

    # Cleanup
    serverA.cancel()
    serverB.cancel()

async def main():
    await test_persistence()
    await test_sync()

if __name__ == "__main__":
    asyncio.run(main())