github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/utils/unittest/cluster.go (about)

     1  package unittest
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"golang.org/x/exp/slices"
     7  
     8  	"github.com/onflow/flow-go/model/flow"
     9  	"github.com/onflow/flow-go/model/flow/factory"
    10  	"github.com/onflow/flow-go/model/flow/filter"
    11  )
    12  
    13  // TransactionForCluster generates a transaction that will be assigned to the
    14  // target cluster ID.
    15  func TransactionForCluster(clusters flow.ClusterList, target flow.IdentitySkeletonList) flow.TransactionBody {
    16  	tx := TransactionBodyFixture()
    17  	return AlterTransactionForCluster(tx, clusters, target, func(*flow.TransactionBody) {})
    18  }
    19  
    20  // AlterTransactionForCluster modifies a transaction nonce until it is assigned
    21  // to the target cluster.
    22  //
    23  // The `after` function is run after each modification to allow for any content
    24  // dependent changes to the transaction (eg. signing it).
    25  func AlterTransactionForCluster(tx flow.TransactionBody, clusters flow.ClusterList, target flow.IdentitySkeletonList, after func(tx *flow.TransactionBody)) flow.TransactionBody {
    26  
    27  	// Bound to avoid infinite loop in case the routing algorithm is broken
    28  	for i := 0; i < 10000; i++ {
    29  		tx.Script = append(tx.Script, '/', '/')
    30  
    31  		if after != nil {
    32  			after(&tx)
    33  		}
    34  		routed, ok := clusters.ByTxID(tx.ID())
    35  		if !ok {
    36  			panic(fmt.Sprintf("unable to find cluster by txID: %x", tx.ID()))
    37  		}
    38  
    39  		if routed.ID() == target.ID() {
    40  			return tx
    41  		}
    42  	}
    43  
    44  	panic(fmt.Sprintf("unable to find transaction for target (%x) with %d clusters", target, len(clusters)))
    45  }
    46  
    47  // ClusterAssignment creates an assignment list with n clusters and with nodes
    48  // evenly distributed among clusters.
    49  func ClusterAssignment(n uint, nodes flow.IdentitySkeletonList) flow.AssignmentList {
    50  
    51  	collectors := nodes.Filter(filter.HasRole[flow.IdentitySkeleton](flow.RoleCollection))
    52  
    53  	// order, so the same list results in the same
    54  	slices.SortFunc(collectors, flow.Canonical[flow.IdentitySkeleton])
    55  
    56  	assignments := make(flow.AssignmentList, n)
    57  	for i, collector := range collectors {
    58  		index := uint(i) % n
    59  		assignments[index] = append(assignments[index], collector.NodeID)
    60  	}
    61  
    62  	return assignments
    63  }
    64  
    65  func ClusterList(n uint, nodes flow.IdentitySkeletonList) flow.ClusterList {
    66  	assignments := ClusterAssignment(n, nodes)
    67  	clusters, err := factory.NewClusterList(assignments, nodes.Filter(filter.HasRole[flow.IdentitySkeleton](flow.RoleCollection)))
    68  	if err != nil {
    69  		panic(err)
    70  	}
    71  
    72  	return clusters
    73  }