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

     1  package unittest
     2  
     3  import (
     4  	"fmt"
     5  	"sort"
     6  
     7  	"github.com/koko1123/flow-go-1/model/flow"
     8  	"github.com/koko1123/flow-go-1/model/flow/factory"
     9  	"github.com/koko1123/flow-go-1/model/flow/filter"
    10  	"github.com/koko1123/flow-go-1/model/flow/order"
    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.IdentityList) 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.IdentityList, 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.Fingerprint() == target.Fingerprint() {
    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.IdentityList) flow.AssignmentList {
    50  
    51  	collectors := nodes.Filter(filter.HasRole(flow.RoleCollection))
    52  
    53  	// order, so the same list results in the same
    54  	sort.Slice(collectors, func(i, j int) bool {
    55  		return order.Canonical(collectors[i], collectors[j])
    56  	})
    57  
    58  	assignments := make(flow.AssignmentList, n)
    59  	for i, collector := range collectors {
    60  		index := uint(i) % n
    61  		assignments[index] = append(assignments[index], collector.NodeID)
    62  	}
    63  
    64  	return assignments
    65  }
    66  
    67  func ClusterList(n uint, nodes flow.IdentityList) flow.ClusterList {
    68  	assignments := ClusterAssignment(n, nodes)
    69  	clusters, err := factory.NewClusterList(assignments, nodes.Filter(filter.HasRole(flow.RoleCollection)))
    70  	if err != nil {
    71  		panic(err)
    72  	}
    73  
    74  	return clusters
    75  }