gitlab.com/gitlab-org/labkit@v1.21.0/correlation/generator_test.go (about)

     1  package correlation
     2  
     3  import (
     4  	"bytes"
     5  	"crypto/rand"
     6  	"strings"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/oklog/ulid/v2"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  )
    14  
    15  func emptyRandomSource() func() {
    16  	oldEntropySource := ulidEntropySource
    17  
    18  	// Use an empty random source, which will lead to ULID generation failure
    19  	ulidEntropySource = &bytes.Buffer{}
    20  
    21  	return func() {
    22  		ulidEntropySource = oldEntropySource
    23  	}
    24  }
    25  
    26  func requireValidRecentULID(require *require.Assertions, got string) {
    27  	uid, err := ulid.Parse(got)
    28  	require.NoError(err, "Expected correlationID to be a valid ULID, got %s", got)
    29  
    30  	utime := ulid.Time(uid.Time())
    31  	diff := time.Since(utime)
    32  	require.True(diff > 0, "Expected ULID to be generated in the past")
    33  	require.True(diff < 1*time.Second, "Expected ULID to be generated with recent timestamp. Timestamp is %v", utime)
    34  }
    35  
    36  func TestRandom(t *testing.T) {
    37  	require := require.New(t)
    38  
    39  	got, err := RandomID()
    40  	require.NoError(err, "Expected no error from RandomID")
    41  
    42  	requireValidRecentULID(require, got)
    43  }
    44  
    45  func TestSafeRandom(t *testing.T) {
    46  	t.Run("is valid", func(t *testing.T) {
    47  		require := require.New(t)
    48  		got := SafeRandomID()
    49  		requireValidRecentULID(require, got)
    50  	})
    51  	t.Run("is random", func(t *testing.T) {
    52  		got1 := SafeRandomID()
    53  		got2 := SafeRandomID()
    54  		require.NotEqual(t, got1, got2)
    55  	})
    56  }
    57  
    58  func TestRandomEntropyFailure(t *testing.T) {
    59  	restore := emptyRandomSource()
    60  	defer restore()
    61  
    62  	require := require.New(t)
    63  
    64  	got, err := RandomID()
    65  
    66  	require.NoError(err, "Expected no error from RandomID")
    67  	require.NotEqual(got, "", "Expected a non-empty string response")
    68  	require.True(strings.HasPrefix(got, "E:"), "Expecting fallback to pseudorandom correlationID")
    69  }
    70  
    71  func TestSafeRandomEntropyFailure(t *testing.T) {
    72  	restore := emptyRandomSource()
    73  	defer restore()
    74  
    75  	require := require.New(t)
    76  
    77  	got := SafeRandomID()
    78  
    79  	require.NotEqual(got, "", "Expected a non-empty string response")
    80  	require.True(strings.HasPrefix(got, "E:"), "Expecting fallback to pseudorandom correlationID")
    81  }
    82  
    83  // TestSafeMonotonicReader tests safeMonotonicReader for data races. It should be ran with -race.
    84  func TestSafeMonotonicReader(t *testing.T) {
    85  	t.Run("MonotonicRead", func(t *testing.T) {
    86  		r := safeMonotonicReader{
    87  			delegate: ulid.Monotonic(rand.Reader, 0),
    88  		}
    89  		go func() {
    90  			d := make([]byte, 100)
    91  			assert.NoError(t, r.MonotonicRead(100, d))
    92  		}()
    93  		go func() {
    94  			d := make([]byte, 100)
    95  			assert.NoError(t, r.MonotonicRead(100, d))
    96  		}()
    97  	})
    98  	t.Run("Read", func(t *testing.T) {
    99  		r := safeMonotonicReader{
   100  			delegate: ulid.Monotonic(rand.Reader, 0),
   101  		}
   102  		go func() {
   103  			d := make([]byte, 100)
   104  			_, err := r.Read(d)
   105  			assert.NoError(t, err)
   106  		}()
   107  		go func() {
   108  			d := make([]byte, 100)
   109  			_, err := r.Read(d)
   110  			assert.NoError(t, err)
   111  		}()
   112  	})
   113  }
   114  
   115  func BenchmarkSafeRandomID(b *testing.B) {
   116  	// run the Fib function b.N times
   117  	for n := 0; n < b.N; n++ {
   118  		SafeRandomID()
   119  	}
   120  }