github.com/decred/dcrlnd@v0.7.6/watchtower/lookout/mock.go (about)

     1  package lookout
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"github.com/decred/dcrd/chaincfg/chainhash"
     8  	"github.com/decred/dcrd/wire"
     9  	"github.com/decred/dcrlnd/chainntnfs"
    10  )
    11  
    12  type MockBackend struct {
    13  	mu sync.RWMutex
    14  
    15  	blocks chan *chainntnfs.BlockEpoch
    16  	epochs map[chainhash.Hash]*wire.MsgBlock
    17  	quit   chan struct{}
    18  }
    19  
    20  func NewMockBackend() *MockBackend {
    21  	return &MockBackend{
    22  		blocks: make(chan *chainntnfs.BlockEpoch),
    23  		epochs: make(map[chainhash.Hash]*wire.MsgBlock),
    24  		quit:   make(chan struct{}),
    25  	}
    26  }
    27  
    28  func (m *MockBackend) RegisterBlockEpochNtfn(*chainntnfs.BlockEpoch) (
    29  	*chainntnfs.BlockEpochEvent, error) {
    30  
    31  	return &chainntnfs.BlockEpochEvent{
    32  		Epochs: m.blocks,
    33  	}, nil
    34  }
    35  
    36  func (m *MockBackend) GetBlock(hash *chainhash.Hash) (*wire.MsgBlock, error) {
    37  	m.mu.RLock()
    38  	defer m.mu.RUnlock()
    39  
    40  	block, ok := m.epochs[*hash]
    41  	if !ok {
    42  		return nil, fmt.Errorf("unknown block for hash %x", hash)
    43  	}
    44  
    45  	return block, nil
    46  }
    47  
    48  func (m *MockBackend) ConnectEpoch(epoch *chainntnfs.BlockEpoch,
    49  	block *wire.MsgBlock) {
    50  
    51  	m.mu.Lock()
    52  	m.epochs[*epoch.Hash] = block
    53  	m.mu.Unlock()
    54  
    55  	select {
    56  	case m.blocks <- epoch:
    57  	case <-m.quit:
    58  	}
    59  }