github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/module/buffer/backend_test.go (about)

     1  package buffer
     2  
     3  import (
     4  	"math/rand"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/suite"
     8  
     9  	"github.com/onflow/flow-go/model/flow"
    10  	"github.com/onflow/flow-go/utils/unittest"
    11  )
    12  
    13  type BackendSuite struct {
    14  	suite.Suite
    15  	backend *backend
    16  }
    17  
    18  func TestBackendSuite(t *testing.T) {
    19  	suite.Run(t, new(BackendSuite))
    20  }
    21  
    22  func (suite *BackendSuite) SetupTest() {
    23  	suite.backend = newBackend()
    24  }
    25  
    26  func (suite *BackendSuite) item() *item {
    27  	parent := unittest.BlockHeaderFixture()
    28  	return suite.itemWithParent(parent)
    29  }
    30  
    31  func (suite *BackendSuite) itemWithParent(parent *flow.Header) *item {
    32  	header := unittest.BlockHeaderWithParentFixture(parent)
    33  	return &item{
    34  		header: flow.Slashable[*flow.Header]{
    35  			OriginID: unittest.IdentifierFixture(),
    36  			Message:  header,
    37  		},
    38  		payload: unittest.IdentifierFixture(),
    39  	}
    40  }
    41  
    42  func (suite *BackendSuite) Add(item *item) {
    43  	suite.backend.add(item.header, item.payload)
    44  }
    45  
    46  func (suite *BackendSuite) TestAdd() {
    47  	expected := suite.item()
    48  	suite.backend.add(expected.header, expected.payload)
    49  
    50  	actual, ok := suite.backend.byID(expected.header.Message.ID())
    51  	suite.Assert().True(ok)
    52  	suite.Assert().Equal(expected, actual)
    53  
    54  	byParent, ok := suite.backend.byParentID(expected.header.Message.ParentID)
    55  	suite.Assert().True(ok)
    56  	suite.Assert().Len(byParent, 1)
    57  	suite.Assert().Equal(expected, byParent[0])
    58  }
    59  
    60  func (suite *BackendSuite) TestChildIndexing() {
    61  
    62  	parent := suite.item()
    63  	child1 := suite.itemWithParent(parent.header.Message)
    64  	child2 := suite.itemWithParent(parent.header.Message)
    65  	grandchild := suite.itemWithParent(child1.header.Message)
    66  	unrelated := suite.item()
    67  
    68  	suite.Add(child1)
    69  	suite.Add(child2)
    70  	suite.Add(grandchild)
    71  	suite.Add(unrelated)
    72  
    73  	suite.Run("retrieve by parent ID", func() {
    74  		byParent, ok := suite.backend.byParentID(parent.header.Message.ID())
    75  		suite.Assert().True(ok)
    76  		// should only include direct children
    77  		suite.Assert().Len(byParent, 2)
    78  		suite.Assert().Contains(byParent, child1)
    79  		suite.Assert().Contains(byParent, child2)
    80  	})
    81  
    82  	suite.Run("drop for parent ID", func() {
    83  		suite.backend.dropForParent(parent.header.Message.ID())
    84  
    85  		// should only drop direct children
    86  		_, exists := suite.backend.byID(child1.header.Message.ID())
    87  		suite.Assert().False(exists)
    88  		_, exists = suite.backend.byID(child2.header.Message.ID())
    89  		suite.Assert().False(exists)
    90  
    91  		// grandchildren should be unaffected
    92  		_, exists = suite.backend.byParentID(child1.header.Message.ID())
    93  		suite.Assert().True(exists)
    94  		_, exists = suite.backend.byID(grandchild.header.Message.ID())
    95  		suite.Assert().True(exists)
    96  
    97  		// nothing else should be affected
    98  		_, exists = suite.backend.byID(unrelated.header.Message.ID())
    99  		suite.Assert().True(exists)
   100  	})
   101  }
   102  
   103  func (suite *BackendSuite) TestPruneByView() {
   104  
   105  	const N = 100 // number of items we're testing with
   106  	items := make([]*item, 0, N)
   107  
   108  	// build a pending buffer
   109  	for i := 0; i < N; i++ {
   110  
   111  		// 10% of the time, add a new unrelated pending header
   112  		if i%10 == 0 {
   113  			item := suite.item()
   114  			suite.Add(item)
   115  			items = append(items, item)
   116  			continue
   117  		}
   118  
   119  		// 90% of the time, build on an existing header
   120  		if i%2 == 1 {
   121  			parent := items[rand.Intn(len(items))]
   122  			item := suite.itemWithParent(parent.header.Message)
   123  			suite.Add(item)
   124  			items = append(items, item)
   125  		}
   126  	}
   127  
   128  	// pick a height to prune that's guaranteed to prune at least one item
   129  	pruneAt := items[rand.Intn(len(items))].header.Message.View
   130  	suite.backend.pruneByView(pruneAt)
   131  
   132  	for _, item := range items {
   133  		view := item.header.Message.View
   134  		id := item.header.Message.ID()
   135  		parentID := item.header.Message.ParentID
   136  
   137  		// check that items below the prune view were removed
   138  		if view <= pruneAt {
   139  			_, exists := suite.backend.byID(id)
   140  			suite.Assert().False(exists)
   141  			_, exists = suite.backend.byParentID(parentID)
   142  			suite.Assert().False(exists)
   143  		}
   144  
   145  		// check that other items were not removed
   146  		if view > item.header.Message.View {
   147  			_, exists := suite.backend.byID(id)
   148  			suite.Assert().True(exists)
   149  			_, exists = suite.backend.byParentID(parentID)
   150  			suite.Assert().True(exists)
   151  		}
   152  	}
   153  }