github.com/onflow/flow-go@v0.33.17/engine/access/state_stream/backend/subscription_test.go (about)

     1  package backend_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/onflow/flow-go/engine/access/state_stream/backend"
    11  
    12  	"github.com/stretchr/testify/assert"
    13  	"github.com/stretchr/testify/require"
    14  
    15  	"github.com/onflow/flow-go/utils/unittest"
    16  )
    17  
    18  // TestSubscription tests that the subscription forwards the data correctly and in order
    19  func TestSubscription_SendReceive(t *testing.T) {
    20  	t.Parallel()
    21  
    22  	ctx := context.Background()
    23  
    24  	sub := backend.NewSubscription(1)
    25  
    26  	assert.NotEmpty(t, sub.ID())
    27  
    28  	messageCount := 20
    29  	messages := []string{}
    30  	for i := 0; i < messageCount; i++ {
    31  		messages = append(messages, fmt.Sprintf("test messages %d", i))
    32  	}
    33  	receivedCount := 0
    34  
    35  	wg := sync.WaitGroup{}
    36  	wg.Add(1)
    37  
    38  	// receive each message and validate it has the expected value
    39  	go func() {
    40  		defer wg.Done()
    41  
    42  		for v := range sub.Channel() {
    43  			assert.Equal(t, messages[receivedCount], v)
    44  			receivedCount++
    45  		}
    46  	}()
    47  
    48  	// send all messages in order
    49  	for _, d := range messages {
    50  		err := sub.Send(ctx, d, 10*time.Millisecond)
    51  		require.NoError(t, err)
    52  	}
    53  	sub.Close()
    54  
    55  	assert.NoError(t, sub.Err())
    56  
    57  	unittest.RequireReturnsBefore(t, wg.Wait, 100*time.Millisecond, "received never finished")
    58  
    59  	assert.Equal(t, messageCount, receivedCount)
    60  }
    61  
    62  // TestSubscription_Failures tests closing and failing subscriptions behaves as expected
    63  func TestSubscription_Failures(t *testing.T) {
    64  	t.Parallel()
    65  
    66  	testErr := fmt.Errorf("test error")
    67  
    68  	// make sure closing a subscription twice does not cause a panic
    69  	t.Run("close only called once", func(t *testing.T) {
    70  		sub := backend.NewSubscription(1)
    71  		sub.Close()
    72  		sub.Close()
    73  
    74  		assert.NoError(t, sub.Err())
    75  	})
    76  
    77  	// make sure failing and closing the same subscription does not cause a panic
    78  	t.Run("close only called once with fail", func(t *testing.T) {
    79  		sub := backend.NewSubscription(1)
    80  		sub.Fail(testErr)
    81  		sub.Close()
    82  
    83  		assert.ErrorIs(t, sub.Err(), testErr)
    84  	})
    85  
    86  	// make sure an error is returned when sending on a closed subscription
    87  	t.Run("send after closed returns an error", func(t *testing.T) {
    88  		sub := backend.NewSubscription(1)
    89  		sub.Fail(testErr)
    90  
    91  		err := sub.Send(context.Background(), "test", 10*time.Millisecond)
    92  		assert.Error(t, err, "expected subscription closed error")
    93  
    94  		assert.ErrorIs(t, sub.Err(), testErr)
    95  	})
    96  }
    97  
    98  // TestHeightBasedSubscription tests that the height based subscription tracks heights correctly
    99  // and forwards the error correctly
   100  func TestHeightBasedSubscription(t *testing.T) {
   101  	t.Parallel()
   102  
   103  	ctx := context.Background()
   104  
   105  	start := uint64(3)
   106  	last := uint64(10)
   107  
   108  	errNoData := fmt.Errorf("no more data")
   109  
   110  	next := start
   111  	getData := func(_ context.Context, height uint64) (interface{}, error) {
   112  		require.Equal(t, next, height)
   113  		if height >= last {
   114  			return nil, errNoData
   115  		}
   116  		next++
   117  		return height, nil
   118  	}
   119  
   120  	// search from [start, last], checking the correct data is returned
   121  	sub := backend.NewHeightBasedSubscription(1, start, getData)
   122  	for i := start; i <= last; i++ {
   123  		data, err := sub.Next(ctx)
   124  		if err != nil {
   125  			// after the last element is returned, next == last
   126  			assert.Equal(t, last, next, "next should be equal to last")
   127  			assert.ErrorIs(t, err, errNoData)
   128  			break
   129  		}
   130  
   131  		require.Equal(t, i, data)
   132  	}
   133  }