github.com/arieschain/arieschain@v0.0.0-20191023063405-37c074544356/p2p/simulations/network_test.go (about)

     1  package simulations
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/quickchainproject/quickchain/p2p/discover"
    10  	"github.com/quickchainproject/quickchain/p2p/simulations/adapters"
    11  )
    12  
    13  // TestNetworkSimulation creates a multi-node simulation network with each node
    14  // connected in a ring topology, checks that all nodes successfully handshake
    15  // with each other and that a snapshot fully represents the desired topology
    16  func TestNetworkSimulation(t *testing.T) {
    17  	// create simulation network with 20 testService nodes
    18  	adapter := adapters.NewSimAdapter(adapters.Services{
    19  		"test": newTestService,
    20  	})
    21  	network := NewNetwork(adapter, &NetworkConfig{
    22  		DefaultService: "test",
    23  	})
    24  	defer network.Shutdown()
    25  	nodeCount := 20
    26  	ids := make([]discover.NodeID, nodeCount)
    27  	for i := 0; i < nodeCount; i++ {
    28  		node, err := network.NewNode()
    29  		if err != nil {
    30  			t.Fatalf("error creating node: %s", err)
    31  		}
    32  		if err := network.Start(node.ID()); err != nil {
    33  			t.Fatalf("error starting node: %s", err)
    34  		}
    35  		ids[i] = node.ID()
    36  	}
    37  
    38  	// perform a check which connects the nodes in a ring (so each node is
    39  	// connected to exactly two peers) and then checks that all nodes
    40  	// performed two handshakes by checking their peerCount
    41  	action := func(_ context.Context) error {
    42  		for i, id := range ids {
    43  			peerID := ids[(i+1)%len(ids)]
    44  			if err := network.Connect(id, peerID); err != nil {
    45  				return err
    46  			}
    47  		}
    48  		return nil
    49  	}
    50  	check := func(ctx context.Context, id discover.NodeID) (bool, error) {
    51  		// check we haven't run out of time
    52  		select {
    53  		case <-ctx.Done():
    54  			return false, ctx.Err()
    55  		default:
    56  		}
    57  
    58  		// get the node
    59  		node := network.GetNode(id)
    60  		if node == nil {
    61  			return false, fmt.Errorf("unknown node: %s", id)
    62  		}
    63  
    64  		// check it has exactly two peers
    65  		client, err := node.Client()
    66  		if err != nil {
    67  			return false, err
    68  		}
    69  		var peerCount int64
    70  		if err := client.CallContext(ctx, &peerCount, "test_peerCount"); err != nil {
    71  			return false, err
    72  		}
    73  		switch {
    74  		case peerCount < 2:
    75  			return false, nil
    76  		case peerCount == 2:
    77  			return true, nil
    78  		default:
    79  			return false, fmt.Errorf("unexpected peerCount: %d", peerCount)
    80  		}
    81  	}
    82  
    83  	timeout := 30 * time.Second
    84  	ctx, cancel := context.WithTimeout(context.Background(), timeout)
    85  	defer cancel()
    86  
    87  	// trigger a check every 100ms
    88  	trigger := make(chan discover.NodeID)
    89  	go triggerChecks(ctx, ids, trigger, 100*time.Millisecond)
    90  
    91  	result := NewSimulation(network).Run(ctx, &Step{
    92  		Action:  action,
    93  		Trigger: trigger,
    94  		Expect: &Expectation{
    95  			Nodes: ids,
    96  			Check: check,
    97  		},
    98  	})
    99  	if result.Error != nil {
   100  		t.Fatalf("simulation failed: %s", result.Error)
   101  	}
   102  
   103  	// take a network snapshot and check it contains the correct topology
   104  	snap, err := network.Snapshot()
   105  	if err != nil {
   106  		t.Fatal(err)
   107  	}
   108  	if len(snap.Nodes) != nodeCount {
   109  		t.Fatalf("expected snapshot to contain %d nodes, got %d", nodeCount, len(snap.Nodes))
   110  	}
   111  	if len(snap.Conns) != nodeCount {
   112  		t.Fatalf("expected snapshot to contain %d connections, got %d", nodeCount, len(snap.Conns))
   113  	}
   114  	for i, id := range ids {
   115  		conn := snap.Conns[i]
   116  		if conn.One != id {
   117  			t.Fatalf("expected conn[%d].One to be %s, got %s", i, id, conn.One)
   118  		}
   119  		peerID := ids[(i+1)%len(ids)]
   120  		if conn.Other != peerID {
   121  			t.Fatalf("expected conn[%d].Other to be %s, got %s", i, peerID, conn.Other)
   122  		}
   123  	}
   124  }
   125  
   126  func triggerChecks(ctx context.Context, ids []discover.NodeID, trigger chan discover.NodeID, interval time.Duration) {
   127  	tick := time.NewTicker(interval)
   128  	defer tick.Stop()
   129  	for {
   130  		select {
   131  		case <-tick.C:
   132  			for _, id := range ids {
   133  				select {
   134  				case trigger <- id:
   135  				case <-ctx.Done():
   136  					return
   137  				}
   138  			}
   139  		case <-ctx.Done():
   140  			return
   141  		}
   142  	}
   143  }