github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/server/internal/notifications/notifications_test.go (about)

     1  package notifications
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/google/fleetspeak/fleetspeak/src/common"
     7  )
     8  
     9  func count(ch <-chan struct{}) int {
    10  	cnt := 0
    11  	for {
    12  		select {
    13  		case _, ok := <-ch:
    14  			if !ok {
    15  				return cnt
    16  			}
    17  			cnt++
    18  		default:
    19  			return cnt
    20  		}
    21  	}
    22  }
    23  
    24  func TestDispatcherDispatches(t *testing.T) {
    25  	id1, err := common.BytesToClientID([]byte{0, 0, 0, 0, 0, 0, 0, 1})
    26  	if err != nil {
    27  		t.Fatal(err)
    28  	}
    29  	id2, err := common.BytesToClientID([]byte{0, 0, 0, 0, 0, 0, 0, 2})
    30  	if err != nil {
    31  		t.Fatal(err)
    32  	}
    33  	id3, err := common.BytesToClientID([]byte{0, 0, 0, 0, 0, 0, 0, 3})
    34  	if err != nil {
    35  		t.Fatal(err)
    36  	}
    37  
    38  	d := NewDispatcher()
    39  	ch1, fin1 := d.Register(id1)
    40  	defer fin1()
    41  	ch2, fin2 := d.Register(id2)
    42  	defer fin2()
    43  
    44  	if cnt := count(ch1); cnt != 0 {
    45  		t.Errorf("Channel should initially be empty, found %d records.", cnt)
    46  	}
    47  
    48  	// Should not block, should put one marker in the channel.
    49  	d.Dispatch(id1)
    50  
    51  	if cnt := count(ch1); cnt != 1 {
    52  		t.Errorf("Dispatch should add one record, but found %d records.", cnt)
    53  	}
    54  	if cnt := count(ch2); cnt != 0 {
    55  		t.Errorf("Unused channel should be empty, but found %d records.", cnt)
    56  	}
    57  
    58  	// Multiple calls should also not block, and should put exactly one marker in
    59  	// the channel.
    60  	d.Dispatch(id1)
    61  	d.Dispatch(id1)
    62  	d.Dispatch(id1)
    63  
    64  	if cnt := count(ch1); cnt != 1 {
    65  		t.Errorf("Dispatch should add one record, but found %d records.", cnt)
    66  	}
    67  	if cnt := count(ch2); cnt != 0 {
    68  		t.Errorf("Unused channel should be empty, but found %d records.", cnt)
    69  	}
    70  
    71  	// Dispatching an unknown id should be a noop.
    72  	d.Dispatch(id3)
    73  	if cnt := count(ch1); cnt != 0 {
    74  		t.Errorf("Unused channel should be empty, but found %d records.", cnt)
    75  	}
    76  	if cnt := count(ch2); cnt != 0 {
    77  		t.Errorf("Unused channel should be empty, but found %d records.", cnt)
    78  	}
    79  }
    80  
    81  func TestDispatcherCloses(t *testing.T) {
    82  	id1, err := common.BytesToClientID([]byte{0, 0, 0, 0, 0, 0, 0, 1})
    83  	if err != nil {
    84  		t.Fatal(err)
    85  	}
    86  
    87  	d := NewDispatcher()
    88  	ch1, fin1 := d.Register(id1)
    89  	ch2, fin2 := d.Register(id1)
    90  
    91  	if _, ok := <-ch1; ok {
    92  		t.Error("ch1 should close with new registration, but read succeeded.")
    93  	}
    94  	// fin1 should be a safe noop at this point.
    95  	fin1()
    96  
    97  	d.Dispatch(id1)
    98  	if cnt := count(ch2); cnt != 1 {
    99  		t.Errorf("Dispatch should add one record, but found %d records.", cnt)
   100  	}
   101  
   102  	fin2()
   103  	if _, ok := <-ch2; ok {
   104  		t.Error("ch2 should close when fin2 is called, but read succeeded.")
   105  	}
   106  }