github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/storage/badger/computation_result_test.go (about)

     1  package badger_test
     2  
     3  import (
     4  	"reflect"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/stretchr/testify/require"
     9  
    10  	"github.com/dgraph-io/badger/v2"
    11  
    12  	"github.com/onflow/flow-go/engine/execution"
    13  	"github.com/onflow/flow-go/engine/execution/testutil"
    14  	bstorage "github.com/onflow/flow-go/storage/badger"
    15  	"github.com/onflow/flow-go/utils/unittest"
    16  )
    17  
    18  func TestUpsertAndRetrieveComputationResult(t *testing.T) {
    19  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    20  		expected := testutil.ComputationResultFixture(t)
    21  		crStorage := bstorage.NewComputationResultUploadStatus(db)
    22  		crId := expected.ExecutableBlock.ID()
    23  
    24  		// True case - upsert
    25  		testUploadStatus := true
    26  		err := crStorage.Upsert(crId, testUploadStatus)
    27  		require.NoError(t, err)
    28  
    29  		actualUploadStatus, err := crStorage.ByID(crId)
    30  		require.NoError(t, err)
    31  
    32  		assert.Equal(t, testUploadStatus, actualUploadStatus)
    33  
    34  		// False case - update
    35  		testUploadStatus = false
    36  		err = crStorage.Upsert(crId, testUploadStatus)
    37  		require.NoError(t, err)
    38  
    39  		actualUploadStatus, err = crStorage.ByID(crId)
    40  		require.NoError(t, err)
    41  
    42  		assert.Equal(t, testUploadStatus, actualUploadStatus)
    43  	})
    44  }
    45  
    46  func TestRemoveComputationResults(t *testing.T) {
    47  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    48  		t.Run("Remove ComputationResult", func(t *testing.T) {
    49  			expected := testutil.ComputationResultFixture(t)
    50  			crId := expected.ExecutableBlock.ID()
    51  			crStorage := bstorage.NewComputationResultUploadStatus(db)
    52  
    53  			testUploadStatus := true
    54  			err := crStorage.Upsert(crId, testUploadStatus)
    55  			require.NoError(t, err)
    56  
    57  			_, err = crStorage.ByID(crId)
    58  			require.NoError(t, err)
    59  
    60  			err = crStorage.Remove(crId)
    61  			require.NoError(t, err)
    62  
    63  			_, err = crStorage.ByID(crId)
    64  			assert.Error(t, err)
    65  		})
    66  	})
    67  }
    68  
    69  func TestListComputationResults(t *testing.T) {
    70  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
    71  		t.Run("List all ComputationResult with given status", func(t *testing.T) {
    72  			expected := [...]*execution.ComputationResult{
    73  				testutil.ComputationResultFixture(t),
    74  				testutil.ComputationResultFixture(t),
    75  			}
    76  			crStorage := bstorage.NewComputationResultUploadStatus(db)
    77  
    78  			// Store a list of ComputationResult instances first
    79  			expectedIDs := make(map[string]bool, 0)
    80  			for _, cr := range expected {
    81  				crId := cr.ExecutableBlock.ID()
    82  				expectedIDs[crId.String()] = true
    83  				err := crStorage.Upsert(crId, true)
    84  				require.NoError(t, err)
    85  			}
    86  			// Add in entries with non-targeted status
    87  			unexpected := [...]*execution.ComputationResult{
    88  				testutil.ComputationResultFixture(t),
    89  				testutil.ComputationResultFixture(t),
    90  			}
    91  			for _, cr := range unexpected {
    92  				crId := cr.ExecutableBlock.ID()
    93  				err := crStorage.Upsert(crId, false)
    94  				require.NoError(t, err)
    95  			}
    96  
    97  			// Get the list of IDs for stored instances
    98  			crIDs, err := crStorage.GetIDsByUploadStatus(true)
    99  			require.NoError(t, err)
   100  
   101  			crIDsStrMap := make(map[string]bool, 0)
   102  			for _, crID := range crIDs {
   103  				crIDsStrMap[crID.String()] = true
   104  			}
   105  
   106  			assert.True(t, reflect.DeepEqual(crIDsStrMap, expectedIDs))
   107  		})
   108  	})
   109  }