ceremonyclient/go-libp2p/p2p/discovery/mocks/mocks.go

package mocks

import (
	"context"
	"sync"
	"time"

	"github.com/libp2p/go-libp2p/core/discovery"
	"github.com/libp2p/go-libp2p/core/host"
	"github.com/libp2p/go-libp2p/core/peer"
)

type clock interface {
	Now() time.Time
}

type MockDiscoveryServer struct {
	mx    sync.Mutex
	db    map[string]map[peer.ID]*discoveryRegistration
	clock clock
}

type discoveryRegistration struct {
	info       peer.AddrInfo
	expiration time.Time
}

func NewDiscoveryServer(clock clock) *MockDiscoveryServer {
	return &MockDiscoveryServer{
		db:    make(map[string]map[peer.ID]*discoveryRegistration),
		clock: clock,
	}
}

func (s *MockDiscoveryServer) Advertise(ns string, info peer.AddrInfo, ttl time.Duration) (time.Duration, error) {
	s.mx.Lock()
	defer s.mx.Unlock()

	peers, ok := s.db[ns]
	if !ok {
		peers = make(map[peer.ID]*discoveryRegistration)
		s.db[ns] = peers
	}
	peers[info.ID] = &discoveryRegistration{info, s.clock.Now().Add(ttl)}
	return ttl, nil
}

func (s *MockDiscoveryServer) FindPeers(ns string, limit int) (<-chan peer.AddrInfo, error) {
	s.mx.Lock()
	defer s.mx.Unlock()

	peers, ok := s.db[ns]
	if !ok || len(peers) == 0 {
		emptyCh := make(chan peer.AddrInfo)
		close(emptyCh)
		return emptyCh, nil
	}

	count := len(peers)
	if limit != 0 && count > limit {
		count = limit
	}

	iterTime := s.clock.Now()
	ch := make(chan peer.AddrInfo, count)
	numSent := 0
	for p, reg := range peers {
		if numSent == count {
			break
		}
		if iterTime.After(reg.expiration) {
			delete(peers, p)
			continue
		}

		numSent++
		ch <- reg.info
	}
	close(ch)

	return ch, nil
}

type MockDiscoveryClient struct {
	host   host.Host
	server *MockDiscoveryServer
}

func NewDiscoveryClient(h host.Host, server *MockDiscoveryServer) *MockDiscoveryClient {
	return &MockDiscoveryClient{
		host:   h,
		server: server,
	}
}

func (d *MockDiscoveryClient) Advertise(ctx context.Context, ns string, opts ...discovery.Option) (time.Duration, error) {
	var options discovery.Options
	err := options.Apply(opts...)
	if err != nil {
		return 0, err
	}

	return d.server.Advertise(ns, *host.InfoFromHost(d.host), options.Ttl)
}

func (d *MockDiscoveryClient) FindPeers(ctx context.Context, ns string, opts ...discovery.Option) (<-chan peer.AddrInfo, error) {
	var options discovery.Options
	err := options.Apply(opts...)
	if err != nil {
		return nil, err
	}

	return d.server.FindPeers(ns, options.Limit)
}