From 16ff4e3ffaf3ce6b62002683902293d9545f7ef4 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Tue, 16 Apr 2019 10:22:04 +0200
Subject: [PATCH] cnct/test: extend mockWitnessBeacon

---
 contractcourt/htlc_timeout_resolver_test.go | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/contractcourt/htlc_timeout_resolver_test.go b/contractcourt/htlc_timeout_resolver_test.go
index c83d66c99..7dcf025ad 100644
--- a/contractcourt/htlc_timeout_resolver_test.go
+++ b/contractcourt/htlc_timeout_resolver_test.go
@@ -29,8 +29,16 @@ func (m *mockSigner) ComputeInputScript(tx *wire.MsgTx,
 
 type mockWitnessBeacon struct {
 	preImageUpdates chan lntypes.Preimage
+	newPreimages    chan []lntypes.Preimage
+	lookupPreimage  map[lntypes.Hash]lntypes.Preimage
+}
 
-	newPreimages chan []lntypes.Preimage
+func newMockWitnessBeacon() *mockWitnessBeacon {
+	return &mockWitnessBeacon{
+		preImageUpdates: make(chan lntypes.Preimage, 1),
+		newPreimages:    make(chan []lntypes.Preimage),
+		lookupPreimage:  make(map[lntypes.Hash]lntypes.Preimage),
+	}
 }
 
 func (m *mockWitnessBeacon) SubscribeUpdates() *WitnessSubscription {
@@ -41,7 +49,11 @@ func (m *mockWitnessBeacon) SubscribeUpdates() *WitnessSubscription {
 }
 
 func (m *mockWitnessBeacon) LookupPreimage(payhash lntypes.Hash) (lntypes.Preimage, bool) {
-	return lntypes.Preimage{}, false
+	preimage, ok := m.lookupPreimage[payhash]
+	if !ok {
+		return lntypes.Preimage{}, false
+	}
+	return preimage, true
 }
 
 func (m *mockWitnessBeacon) AddPreimages(preimages ...lntypes.Preimage) error {
@@ -190,10 +202,7 @@ func TestHtlcTimeoutResolver(t *testing.T) {
 		spendChan: make(chan *chainntnfs.SpendDetail),
 		confChan:  make(chan *chainntnfs.TxConfirmation),
 	}
-	witnessBeacon := &mockWitnessBeacon{
-		preImageUpdates: make(chan lntypes.Preimage, 1),
-		newPreimages:    make(chan []lntypes.Preimage),
-	}
+	witnessBeacon := newMockWitnessBeacon()
 
 	for _, testCase := range testCases {
 		t.Logf("Running test case: %v", testCase.name)