github.com/anycable/anycable-go@v1.5.1/hub/gate_test.go (about)

     1  package hub
     2  
     3  import (
     4  	"context"
     5  	"log/slog"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/anycable/anycable-go/common"
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  func TestUnsubscribe(t *testing.T) {
    15  	ctx, cancel := context.WithCancel(context.Background())
    16  	defer cancel()
    17  
    18  	gate := NewGate(ctx, slog.Default())
    19  
    20  	session := NewMockSession("123")
    21  
    22  	gate.Subscribe(session, "test", "test_channel")
    23  
    24  	assert.Equal(t, 1, gate.Size())
    25  
    26  	gate.Unsubscribe(session, "test", "test_channel")
    27  
    28  	assert.Equal(t, 0, gate.Size())
    29  
    30  	assert.Equal(t, 0, len(gate.streams))
    31  	assert.Equal(t, 0, len(gate.sessionsStreams))
    32  }
    33  
    34  func TestShutdown(t *testing.T) {
    35  	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
    36  	defer cancel()
    37  
    38  	gate := NewGate(ctx, slog.Default())
    39  
    40  	session := NewMockSession("123")
    41  
    42  	gate.Subscribe(session, "test", "test_channel")
    43  
    44  	gate.Broadcast(&common.StreamMessage{Stream: "test", Data: "1"})
    45  
    46  	<-ctx.Done()
    47  
    48  	gate.Broadcast(&common.StreamMessage{Stream: "test", Data: "2"})
    49  
    50  	// Ignore first message if any, we want to make sure the second one is not received
    51  	session.Read() // nolint:errcheck
    52  
    53  	msg, err := session.Read()
    54  
    55  	require.Error(t, err, "expected not to receive messages when context is done, but received: %s", msg)
    56  }