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

     1  package unittest
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/libp2p/go-libp2p/core/peer"
     7  
     8  	"github.com/onflow/flow-go/model/flow"
     9  	"github.com/onflow/flow-go/network/p2p/keyutils"
    10  )
    11  
    12  // UpdatableIDProvider implements an IdentityProvider which can be manually updated by setting
    13  // the IdentityList to a new value.
    14  // It also implements an IdentifierProvider which provides the identifiers of the IdentityList.
    15  // This is mainly used to simulate epoch transitions in tests.
    16  type UpdatableIDProvider struct {
    17  	mu         sync.RWMutex
    18  	identities flow.IdentityList
    19  }
    20  
    21  func NewUpdatableIDProvider(identities flow.IdentityList) *UpdatableIDProvider {
    22  	return &UpdatableIDProvider{identities: identities}
    23  }
    24  
    25  // SetIdentities updates the IdentityList returned by this provider.
    26  func (p *UpdatableIDProvider) SetIdentities(identities flow.IdentityList) {
    27  	p.mu.Lock()
    28  	defer p.mu.Unlock()
    29  	p.identities = identities
    30  }
    31  
    32  func (p *UpdatableIDProvider) Identities(filter flow.IdentityFilter[flow.Identity]) flow.IdentityList {
    33  	p.mu.RLock()
    34  	defer p.mu.RUnlock()
    35  	return p.identities.Filter(filter)
    36  }
    37  
    38  func (p *UpdatableIDProvider) Identifiers() flow.IdentifierList {
    39  	p.mu.RLock()
    40  	defer p.mu.RUnlock()
    41  	return p.identities.NodeIDs()
    42  }
    43  
    44  func (p *UpdatableIDProvider) ByNodeID(flowID flow.Identifier) (*flow.Identity, bool) {
    45  	p.mu.RLock()
    46  	defer p.mu.RUnlock()
    47  
    48  	for _, v := range p.identities {
    49  		if v.ID() == flowID {
    50  			return v, true
    51  		}
    52  	}
    53  	return nil, false
    54  }
    55  
    56  func (p *UpdatableIDProvider) ByPeerID(peerID peer.ID) (*flow.Identity, bool) {
    57  	p.mu.RLock()
    58  	defer p.mu.RUnlock()
    59  
    60  	for _, v := range p.identities {
    61  		if id, err := keyutils.PeerIDFromFlowPublicKey(v.NetworkPubKey); err == nil {
    62  			if id == peerID {
    63  				return v, true
    64  			}
    65  		}
    66  
    67  	}
    68  	return nil, false
    69  
    70  }