github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/engine/access/subscription/subscription_test.go (about)

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