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 }