test_node.py

import threading
import time
import sys
from messaging import MessagingModule
from protocol import *
from discovery import DiscoveryModule
from bully_algorithm import BullyAlgorithm, PeerInfo

def main():
    if len(sys.argv) < 3:
        print("Usage: python test_node.py <port> <node_id>")
        return
        
    # Listen on all interfaces to allow connections from LAN IP
    my_ip = "0.0.0.0" 
    my_port = int(sys.argv[1])
    my_id = int(sys.argv[2]) 
    
    print(f"Starting Node {my_id} on {my_ip}:{my_port}")

    # 1. Initialize messaging module
    messenger = MessagingModule(my_ip, my_port, my_id)
    
    # 2. Initialize Bully Algorithm (initially empty peers)
    bully = BullyAlgorithm(my_id, messenger, {})
    messenger.logic_handler = bully
    
    # 3. Define callback for Discovery
    def on_peers_update(discovered_peers):
        # Convert DiscoveredPeer to PeerInfo
        bully_peers = {}
        for pid, p in discovered_peers.items():
            bully_peers[pid] = PeerInfo(pid, p.host, p.port)
        
        bully.update_peers(bully_peers)

    # 4. Initialize Discovery Module
    discovery = DiscoveryModule(my_id, my_port, on_peer_changed=on_peers_update)
    
    # 5. Start threads
    t_net = threading.Thread(target=messenger.start_listening, daemon=True)
    t_net.start()
    
    discovery.start()

    # 6. Start status reporting thread
    def report_status():
        while True:
            time.sleep(5)
            print("\n" + "="*40)
            print(f"📊 [Status Report] Node {my_id}")
            print(f"   - Leader: {bully.leader_id}")
            print(f"   - State: {'Election' if bully.in_election else 'Normal'}")
            print(f"   - Peers ({len(bully.peers)}): {list(bully.peers.keys())}")
            
            # Simulate Heartbeat logic (just for reporting)
            if bully.leader_id and bully.leader_id != my_id:
                print(f"   - ❤️ Sending Heartbeat to Leader {bully.leader_id}...")
                # In a real implementation, you would call messenger.send_message here
            elif bully.leader_id == my_id:
                print(f"   - 👑 I am Leader. Waiting for heartbeats.")
            print("="*40 + "\n")

    threading.Thread(target=report_status, daemon=True).start()
    
    # User input loop
    print("Commands:")
    print("  e: Start Election")
    print("  s: Show Status")
    print("  p: Show Peers")
    print("  q: Quit")
    
    while True:
        try:
            cmd = input(f"Node {my_id} > ").strip().lower()
            if not cmd: continue
            
            if cmd == 'e':
                bully.start_election()
            elif cmd == 's':
                print(f"Current Leader: {bully.leader_id}")
                print(f"In Election: {bully.in_election}")
            elif cmd == 'p':
                print(f"Peers: {list(bully.peers.keys())}")
            elif cmd == 'q':
                discovery.stop()
                break
        except KeyboardInterrupt:
            break
            
if __name__ == "__main__":
    main()