github.com/ethersphere/bee/v2@v2.2.0/pkg/discovery/mock/mock.go (about)

     1  // Copyright 2020 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package mock
     6  
     7  import (
     8  	"context"
     9  	"sync"
    10  
    11  	"github.com/ethersphere/bee/v2/pkg/swarm"
    12  )
    13  
    14  type Discovery struct {
    15  	mtx           sync.Mutex
    16  	ctr           int //how many ops
    17  	records       map[string][]swarm.Address
    18  	broadcastFunc func(context.Context, swarm.Address, ...swarm.Address) error
    19  }
    20  
    21  type Option interface {
    22  	apply(*Discovery)
    23  }
    24  type optionFunc func(*Discovery)
    25  
    26  func (f optionFunc) apply(r *Discovery) { f(r) }
    27  
    28  func WithBroadcastPeers(f func(context.Context, swarm.Address, ...swarm.Address) error) optionFunc {
    29  	return optionFunc(func(r *Discovery) {
    30  		r.broadcastFunc = f
    31  	})
    32  }
    33  
    34  func NewDiscovery(opts ...Option) *Discovery {
    35  	d := &Discovery{
    36  		records: make(map[string][]swarm.Address),
    37  	}
    38  	for _, opt := range opts {
    39  		opt.apply(d)
    40  	}
    41  	return d
    42  }
    43  
    44  func (d *Discovery) BroadcastPeers(ctx context.Context, addressee swarm.Address, peers ...swarm.Address) error {
    45  	if d.broadcastFunc != nil {
    46  		return d.broadcastFunc(ctx, addressee, peers...)
    47  	}
    48  	for _, peer := range peers {
    49  		d.mtx.Lock()
    50  		d.records[addressee.String()] = append(d.records[addressee.String()], peer)
    51  		d.mtx.Unlock()
    52  	}
    53  	d.mtx.Lock()
    54  	d.ctr++
    55  	d.mtx.Unlock()
    56  	return nil
    57  }
    58  
    59  func (d *Discovery) Broadcasts() int {
    60  	d.mtx.Lock()
    61  	defer d.mtx.Unlock()
    62  	return d.ctr
    63  }
    64  
    65  func (d *Discovery) AddresseeRecords(addressee swarm.Address) (peers []swarm.Address, exists bool) {
    66  	d.mtx.Lock()
    67  	defer d.mtx.Unlock()
    68  	peers, exists = d.records[addressee.String()]
    69  	return
    70  }
    71  
    72  func (d *Discovery) Reset() {
    73  	d.mtx.Lock()
    74  	defer d.mtx.Unlock()
    75  	d.ctr = 0
    76  	d.records = make(map[string][]swarm.Address)
    77  }