github.com/jincm/wesharechain@v0.0.0-20210122032815-1537409ce26a/chain/swarm/network/simulation/events_test.go (about)

     1  // Copyright 2018 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package simulation
    18  
    19  import (
    20  	"context"
    21  	"sync"
    22  	"testing"
    23  	"time"
    24  )
    25  
    26  // TestPeerEvents creates simulation, adds two nodes,
    27  // register for peer events, connects nodes in a chain
    28  // and waits for the number of connection events to
    29  // be received.
    30  func TestPeerEvents(t *testing.T) {
    31  	sim := New(noopServiceFuncMap)
    32  	defer sim.Close()
    33  
    34  	_, err := sim.AddNodes(2)
    35  	if err != nil {
    36  		t.Fatal(err)
    37  	}
    38  
    39  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    40  	defer cancel()
    41  	events := sim.PeerEvents(ctx, sim.NodeIDs())
    42  
    43  	// two nodes -> two connection events
    44  	expectedEventCount := 2
    45  
    46  	var wg sync.WaitGroup
    47  	wg.Add(expectedEventCount)
    48  
    49  	go func() {
    50  		for e := range events {
    51  			if e.Error != nil {
    52  				if e.Error == context.Canceled {
    53  					return
    54  				}
    55  				t.Error(e.Error)
    56  				continue
    57  			}
    58  			wg.Done()
    59  		}
    60  	}()
    61  
    62  	err = sim.Net.ConnectNodesChain(sim.NodeIDs())
    63  	if err != nil {
    64  		t.Fatal(err)
    65  	}
    66  
    67  	wg.Wait()
    68  }
    69  
    70  func TestPeerEventsTimeout(t *testing.T) {
    71  	sim := New(noopServiceFuncMap)
    72  	defer sim.Close()
    73  
    74  	_, err := sim.AddNodes(2)
    75  	if err != nil {
    76  		t.Fatal(err)
    77  	}
    78  
    79  	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
    80  	defer cancel()
    81  	events := sim.PeerEvents(ctx, sim.NodeIDs())
    82  
    83  	done := make(chan struct{})
    84  	errC := make(chan error)
    85  	go func() {
    86  		for e := range events {
    87  			if e.Error == context.Canceled {
    88  				return
    89  			}
    90  			if e.Error == context.DeadlineExceeded {
    91  				close(done)
    92  				return
    93  			} else {
    94  				errC <- e.Error
    95  			}
    96  		}
    97  	}()
    98  
    99  	select {
   100  	case <-time.After(time.Second):
   101  		t.Fatal("no context deadline received")
   102  	case err := <-errC:
   103  		t.Fatal(err)
   104  	case <-done:
   105  		// all good, context deadline detected
   106  	}
   107  }