github.com/neatlab/neatio@v1.7.3-0.20220425043230-d903e92fcc75/network/p2p/simulations/network_test.go (about)

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