github.com/martinohmann/rfoutlet@v1.2.1-0.20220707195255-8a66aa411105/internal/websocket/hub_test.go (about)

     1  package websocket
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func TestHub(t *testing.T) {
    12  	h := NewHub()
    13  
    14  	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    15  	defer cancel()
    16  
    17  	go h.Run(ctx.Done())
    18  
    19  	c1 := newClient(h, nil, nil)
    20  	c2 := newClient(h, nil, nil)
    21  
    22  	go func() {
    23  		defer cancel()
    24  
    25  		h.register <- c1
    26  		h.register <- c2
    27  
    28  		h.Send(c1, []byte{0x1})
    29  		assert.Equal(t, []byte{0x1}, <-c1.send)
    30  
    31  		h.Send(c2, []byte{0x2})
    32  		assert.Equal(t, []byte{0x2}, <-c2.send)
    33  
    34  		h.Broadcast([]byte{0x3})
    35  		assert.Equal(t, []byte{0x3}, <-c1.send)
    36  		assert.Equal(t, []byte{0x3}, <-c2.send)
    37  
    38  		h.unregister <- c1
    39  
    40  		h.Broadcast([]byte{0x4})
    41  		assert.Equal(t, []byte{0x4}, <-c2.send)
    42  	}()
    43  
    44  	<-ctx.Done()
    45  
    46  	if err := ctx.Err(); err == context.DeadlineExceeded {
    47  		t.Fatal(err)
    48  	}
    49  }