github.com/koko1123/flow-go-1@v0.29.6/utils/unittest/cluster_state_checker.go (about)

     1  package unittest
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/assert"
     7  
     8  	"github.com/koko1123/flow-go-1/model/flow"
     9  	"github.com/koko1123/flow-go-1/state/cluster"
    10  )
    11  
    12  const dontCheck = -1
    13  
    14  // ClusterStateChecker is a test utility for checking cluster state. First, prepare it
    15  // with expectations about the state, using `Expect*` functions, then use
    16  // `Check` to assert the expectations.
    17  //
    18  // Duplicates are checked automatically without setting any expectations.
    19  type ClusterStateChecker struct {
    20  	state cluster.State
    21  
    22  	// expectations about state
    23  	count    int
    24  	contains []flow.Identifier
    25  	omits    []flow.Identifier
    26  }
    27  
    28  // NewClusterStateChecker returns a state checker for the given state.
    29  func NewClusterStateChecker(state cluster.State) *ClusterStateChecker {
    30  	checker := &ClusterStateChecker{
    31  		state:    state,
    32  		count:    dontCheck,
    33  		contains: nil,
    34  		omits:    nil,
    35  	}
    36  	return checker
    37  }
    38  
    39  // ExpectTxCount adds an expectation for the total count of transactions in
    40  // the cluster state.
    41  func (checker *ClusterStateChecker) ExpectTxCount(n int) *ClusterStateChecker {
    42  	checker.count = n
    43  	return checker
    44  }
    45  
    46  // ExpectContainsTx adds an expectation that the given transaction exists in
    47  // the cluster state.
    48  func (checker *ClusterStateChecker) ExpectContainsTx(txIDs ...flow.Identifier) *ClusterStateChecker {
    49  	checker.contains = append(checker.contains, txIDs...)
    50  	return checker
    51  }
    52  
    53  // ExpectOmitsTx adds an expectation that the given  transaction does not exist
    54  // in the cluster state.
    55  func (checker *ClusterStateChecker) ExpectOmitsTx(txIDs ...flow.Identifier) *ClusterStateChecker {
    56  	checker.omits = append(checker.omits, txIDs...)
    57  	return checker
    58  }
    59  
    60  // Assert checks all assertions against the cluster state. If any assertions
    61  // fail, the test will fail.
    62  func (checker *ClusterStateChecker) Assert(t *testing.T) {
    63  
    64  	// start at the state head
    65  	head, err := checker.state.Final().Head()
    66  	assert.Nil(t, err)
    67  
    68  	// track properties of the state we will later compare against expectations
    69  	var (
    70  		count        = 0                                  // total number of transactions
    71  		transactions = make(map[flow.Identifier]struct{}) // unique transactions
    72  		dupes        []flow.Identifier                    // duplicate transactions
    73  	)
    74  
    75  	// walk the chain state from head to genesis
    76  	for head.Height > 0 {
    77  		collection, err := checker.state.AtBlockID(head.ID()).Collection()
    78  		assert.Nil(t, err)
    79  
    80  		head, err = checker.state.AtBlockID(head.ParentID).Head()
    81  		assert.Nil(t, err)
    82  
    83  		if collection.Len() == 0 {
    84  			continue
    85  		}
    86  
    87  		for _, txID := range collection.Light().Transactions {
    88  			count++
    89  
    90  			_, isDupe := transactions[txID]
    91  			if isDupe {
    92  				dupes = append(dupes, txID)
    93  			}
    94  
    95  			transactions[txID] = struct{}{}
    96  		}
    97  	}
    98  
    99  	// ensure there are no duplicates
   100  	if !assert.Len(t, dupes, 0) {
   101  		t.Log("found duplicates: ", dupes)
   102  	}
   103  
   104  	// check that all manually set expectations are true
   105  	if checker.count != dontCheck {
   106  		assert.Equal(t, checker.count, count, "unexpected total number of transactions")
   107  	}
   108  
   109  	for _, txID := range checker.contains {
   110  		_, exists := transactions[txID]
   111  		assert.True(t, exists, "missing expected transaction: %x", txID)
   112  	}
   113  
   114  	for _, txID := range checker.omits {
   115  		_, exists := transactions[txID]
   116  		assert.False(t, exists, "found unexpected transaction: %x", txID)
   117  	}
   118  }