test_results.py

#!/usr/bin/env python3
"""
Test script to verify Shadow simulation results for GossipSub by reading log files.
Verifies that all subscribers received messages from all publishers.

Usage: python3 test_results.py [node_count] [nodes_to_publish] [nodes_to_subscribe]
"""

import sys
import os
import glob
from typing import Dict, List
import json

def parse_shadow_logs(node_count: int, subscribed_node_ids: List[int]) -> Dict[int, int]:
    """
    Parse Shadow log files to extract message reception counts for subscribers.

    Returns:
        - received_counts: Dict mapping node_id to number of messages received
    """
    received_counts = {}

    # Check subscribers (nodes 0 to nodes_to_subscribe-1) since they are the ones subscribed to the topic
    for node_id in subscribed_node_ids:
        node_dir = f"shadow.data/hosts/node{node_id}"
        received_count = 0

        if os.path.exists(node_dir):
            # Find stderr files (Shadow logs are in stderr for Go programs)
            stderr_files = glob.glob(f"{node_dir}/*.stderr")

            if stderr_files:
                try:
                    # Use the first stderr file found
                    log_file = stderr_files[0]

                    with open(log_file, 'r') as f:
                        for line in f:
                            # Look for "Received message" log entries
                            if "Received message" in line:
                                received_count += 1

                except Exception as e:
                    print(f"Warning: Could not read {log_file}: {e}")
            else:
                print(f"Warning: No stderr files found in {node_dir}")
        else:
            print(f"Warning: Node directory not found: {node_dir}")

        received_counts[node_id] = received_count

    return received_counts

def test_message_delivery(node_count: int, topology_file: str) -> bool:
    """Test that all subscribers received messages from all publishers."""
    topology = None
    with open(topology_file, 'r') as f:
        topology = json.load(f)

    mesh_nodes = len(topology["mesh_node_ids"])
    mesh_attester_nodes = len(topology["mesh_attester_node_ids"])
    non_mesh_attester_nodes = len(topology["non_mesh_attester_node_ids"])

    nodes_to_subscribe = mesh_nodes + mesh_attester_nodes
    nodes_to_publish = non_mesh_attester_nodes + mesh_attester_nodes

    subscribed_node_ids = topology["mesh_node_ids"] + topology["mesh_attester_node_ids"]
    publisher_node_ids = topology["non_mesh_attester_node_ids"] + topology["mesh_attester_node_ids"]

    print(f"Shadow GossipSub Simulation Test Results")
    print("=" * 60)
    print(f"Total nodes: {node_count}, Mesh Nodes: {mesh_nodes}, Mesh Attesters: {mesh_attester_nodes}, Non Mesh Attesters: {non_mesh_attester_nodes}")
    print(f"Testing message delivery to {nodes_to_subscribe} subscribers from {nodes_to_publish} publishers...")

    received_counts = parse_shadow_logs(node_count, subscribed_node_ids)

    all_passed = True
    total_received = 0
    expected_messages = nodes_to_publish  # Each subscriber should receive messages from all publishers

    print(f"\nMessage Delivery to Subscribers (Expected: {expected_messages} messages per subscriber)")
    print("-" * 60)

    for node_id in subscribed_node_ids:
        received = received_counts.get(node_id, 0)
        total_received += received

        if received == expected_messages:
            status = "✓ PASS"
        else:
            status = "✗ FAIL"
            all_passed = False

        print(f"Subscriber {node_id:2d}: {received}/{expected_messages} messages {status}")

    print("-" * 60)
    print(f"Total messages received: {total_received}/{nodes_to_subscribe * expected_messages}")

    print()
    if all_passed:
        print("✓ ALL TESTS PASSED: All subscribers received all published messages")
        return True
    else:
        failed_subscribers = [i for i in subscribed_node_ids if received_counts.get(i, 0) != expected_messages]
        print(f"✗ TEST FAILED: {len(failed_subscribers)} subscriber(s) failed to receive all published messages")
        print(f"Failed subscribers: {failed_subscribers}")
        return False

def check_shadow_data_exists() -> bool:
    """Check if Shadow simulation data directory exists."""
    if not os.path.exists("shadow.data"):
        print("Error: shadow.data directory not found. Run the simulation first.")
        print("  Run: make all")
        return False

    if not os.path.exists("shadow.data/hosts"):
        print("Error: shadow.data/hosts directory not found. Simulation may have failed.")
        return False

    return True

def main():
    """Main function."""
    if len(sys.argv) < 2:
        print("Usage: python3 test_results.py [node_count] [topology_file]")
        print("Example: python3 test_results.py 10 topology.json")
        sys.exit(1)

    try:
        node_count = int(sys.argv[1])
        topology_file = sys.argv[2]
    except ValueError:
        print("Error: node_count and topology_file must be valid")
        sys.exit(1)

    if not check_shadow_data_exists():
        sys.exit(1)

    success = test_message_delivery(node_count, topology_file)
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()