github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/utils/concurrentqueue/concurrentqueue_test.go (about)

     1  package concurrentqueue
     2  
     3  import (
     4  	"sync"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/require"
     8  )
     9  
    10  type testEntry struct {
    11  	workerIndex int
    12  	value       int
    13  }
    14  
    15  // TestSafeQueue checks if `Push` and `Pop` operations work in concurrent environment
    16  func TestSafeQueue(t *testing.T) {
    17  	var q ConcurrentQueue
    18  	workers := 1000
    19  	samples := 100
    20  	wg := sync.WaitGroup{}
    21  	wg.Add(workers)
    22  	for i := 0; i < workers; i++ {
    23  		go func(worker int) {
    24  			for i := 1; i <= samples; i++ {
    25  				q.Push(testEntry{
    26  					workerIndex: worker,
    27  					value:       i,
    28  				})
    29  			}
    30  			wg.Done()
    31  		}(i)
    32  	}
    33  
    34  	wg.Wait()
    35  
    36  	r := make([]int, workers)
    37  
    38  	// accumulate all values
    39  	for {
    40  		v, found := q.Pop()
    41  		if !found {
    42  			break
    43  		}
    44  		entry := v.(testEntry)
    45  		r[entry.workerIndex] += entry.value
    46  	}
    47  	// (a1 + aN) * N / 2
    48  	expected := (samples + 1) * samples / 2
    49  	for _, acc := range r {
    50  		require.Equal(t, expected, acc)
    51  	}
    52  }
    53  
    54  // TestPopBatch checks correctness of `PopBatch`.
    55  func TestPopBatch(t *testing.T) {
    56  	var q ConcurrentQueue
    57  	samples := 100
    58  	for i := 1; i <= samples; i++ {
    59  		q.Push(i)
    60  	}
    61  
    62  	// (a1 + aN) * N / 2
    63  	expected := (samples + 1) * samples / 2
    64  	actual := 0
    65  	for {
    66  		v, found := q.PopBatch(11)
    67  		if !found {
    68  			break
    69  		}
    70  		for _, item := range v {
    71  			actual += item.(int)
    72  		}
    73  	}
    74  
    75  	require.Equal(t, 0, q.Len())
    76  	require.Equal(t, expected, actual)
    77  }