From 07336712ad311acdec4e88c52aef2aac5461678f Mon Sep 17 00:00:00 2001
From: William Banfield <wbanfield@gmail.com>
Date: Thu, 20 Jan 2022 17:07:46 -0500
Subject: [PATCH] context add to randLightBlock

---
 light/store/db/db_test.go | 27 ++++++++++++++++++++-------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/light/store/db/db_test.go b/light/store/db/db_test.go
index 300456833..aa9539f99 100644
--- a/light/store/db/db_test.go
+++ b/light/store/db/db_test.go
@@ -1,6 +1,7 @@
 package db
 
 import (
+	"context"
 	"sync"
 	"testing"
 	"time"
@@ -19,6 +20,8 @@ import (
 
 func TestLast_FirstLightBlockHeight(t *testing.T) {
 	dbStore := New(dbm.NewMemDB())
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	// Empty store
 	height, err := dbStore.LastLightBlockHeight()
@@ -30,7 +33,7 @@ func TestLast_FirstLightBlockHeight(t *testing.T) {
 	assert.EqualValues(t, -1, height)
 
 	// 1 key
-	err = dbStore.SaveLightBlock(randLightBlock(t, int64(1)))
+	err = dbStore.SaveLightBlock(randLightBlock(t, ctx, int64(1)))
 	require.NoError(t, err)
 
 	height, err = dbStore.LastLightBlockHeight()
@@ -44,6 +47,8 @@ func TestLast_FirstLightBlockHeight(t *testing.T) {
 
 func Test_SaveLightBlock(t *testing.T) {
 	dbStore := New(dbm.NewMemDB())
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	// Empty store
 	h, err := dbStore.LightBlock(1)
@@ -51,7 +56,7 @@ func Test_SaveLightBlock(t *testing.T) {
 	assert.Nil(t, h)
 
 	// 1 key
-	err = dbStore.SaveLightBlock(randLightBlock(t, 1))
+	err = dbStore.SaveLightBlock(randLightBlock(t, ctx, 1))
 	require.NoError(t, err)
 
 	size := dbStore.Size()
@@ -74,13 +79,15 @@ func Test_SaveLightBlock(t *testing.T) {
 
 func Test_LightBlockBefore(t *testing.T) {
 	dbStore := New(dbm.NewMemDB())
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	assert.Panics(t, func() {
 		_, _ = dbStore.LightBlockBefore(0)
 		_, _ = dbStore.LightBlockBefore(100)
 	})
 
-	err := dbStore.SaveLightBlock(randLightBlock(t, int64(2)))
+	err := dbStore.SaveLightBlock(randLightBlock(t, ctx, int64(2)))
 	require.NoError(t, err)
 
 	h, err := dbStore.LightBlockBefore(3)
@@ -95,6 +102,8 @@ func Test_LightBlockBefore(t *testing.T) {
 
 func Test_Prune(t *testing.T) {
 	dbStore := New(dbm.NewMemDB())
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	// Empty store
 	assert.EqualValues(t, 0, dbStore.Size())
@@ -102,7 +111,7 @@ func Test_Prune(t *testing.T) {
 	require.NoError(t, err)
 
 	// One header
-	err = dbStore.SaveLightBlock(randLightBlock(t, 2))
+	err = dbStore.SaveLightBlock(randLightBlock(t, ctx, 2))
 	require.NoError(t, err)
 
 	assert.EqualValues(t, 1, dbStore.Size())
@@ -117,7 +126,7 @@ func Test_Prune(t *testing.T) {
 
 	// Multiple headers
 	for i := 1; i <= 10; i++ {
-		err = dbStore.SaveLightBlock(randLightBlock(t, int64(i)))
+		err = dbStore.SaveLightBlock(randLightBlock(t, ctx, int64(i)))
 		require.NoError(t, err)
 	}
 
@@ -133,13 +142,16 @@ func Test_Prune(t *testing.T) {
 func Test_Concurrency(t *testing.T) {
 	dbStore := New(dbm.NewMemDB())
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	var wg sync.WaitGroup
 	for i := 1; i <= 100; i++ {
 		wg.Add(1)
 		go func(i int64) {
 			defer wg.Done()
 
-			err := dbStore.SaveLightBlock(randLightBlock(t, i))
+			err := dbStore.SaveLightBlock(randLightBlock(t, ctx, i))
 			require.NoError(t, err)
 
 			_, err = dbStore.LightBlock(i)
@@ -182,7 +194,8 @@ func Test_Concurrency(t *testing.T) {
 	wg.Wait()
 }
 
-func randLightBlock(t *testing.T, height int64) *types.LightBlock {
+func randLightBlock(t *testing.T, ctx context.Context, height int64) *types.LightBlock {
+	t.Helper()
 	vals, _ := factory.ValidatorSet(t, ctx, 2, 1)
 	return &types.LightBlock{
 		SignedHeader: &types.SignedHeader{