From 6f06c304219fc6c8351c553238ca350ce7f062fc Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Tue, 19 Mar 2019 19:22:47 -0700
Subject: [PATCH] mock: support late registration of spend ntfns

---
 mock.go | 46 ++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 38 insertions(+), 8 deletions(-)

diff --git a/mock.go b/mock.go
index 1eccbea6..6c1f3fb6 100644
--- a/mock.go
+++ b/mock.go
@@ -123,6 +123,7 @@ func (m *mockNotfier) RegisterSpendNtfn(outpoint *wire.OutPoint, _ []byte,
 type mockSpendNotifier struct {
 	*mockNotfier
 	spendMap map[wire.OutPoint][]chan *chainntnfs.SpendDetail
+	spends   map[wire.OutPoint]*chainntnfs.SpendDetail
 	mtx      sync.Mutex
 }
 
@@ -132,6 +133,7 @@ func makeMockSpendNotifier() *mockSpendNotifier {
 			confChannel: make(chan *chainntnfs.TxConfirmation),
 		},
 		spendMap: make(map[wire.OutPoint][]chan *chainntnfs.SpendDetail),
+		spends:   make(map[wire.OutPoint]*chainntnfs.SpendDetail),
 	}
 }
 
@@ -140,8 +142,22 @@ func (m *mockSpendNotifier) RegisterSpendNtfn(outpoint *wire.OutPoint,
 	m.mtx.Lock()
 	defer m.mtx.Unlock()
 
-	spendChan := make(chan *chainntnfs.SpendDetail)
-	m.spendMap[*outpoint] = append(m.spendMap[*outpoint], spendChan)
+	spendChan := make(chan *chainntnfs.SpendDetail, 1)
+	if detail, ok := m.spends[*outpoint]; ok {
+		// Deliver spend immediately if details are already known.
+		spendChan <- &chainntnfs.SpendDetail{
+			SpentOutPoint:     detail.SpentOutPoint,
+			SpendingHeight:    detail.SpendingHeight,
+			SpendingTx:        detail.SpendingTx,
+			SpenderTxHash:     detail.SpenderTxHash,
+			SpenderInputIndex: detail.SpenderInputIndex,
+		}
+	} else {
+		// Otherwise, queue the notification for delivery if the spend
+		// is ever received.
+		m.spendMap[*outpoint] = append(m.spendMap[*outpoint], spendChan)
+	}
+
 	return &chainntnfs.SpendEvent{
 		Spend: spendChan,
 		Cancel: func() {
@@ -156,16 +172,30 @@ func (m *mockSpendNotifier) Spend(outpoint *wire.OutPoint, height int32,
 	m.mtx.Lock()
 	defer m.mtx.Unlock()
 
+	txnHash := txn.TxHash()
+	details := &chainntnfs.SpendDetail{
+		SpentOutPoint:     outpoint,
+		SpendingHeight:    height,
+		SpendingTx:        txn,
+		SpenderTxHash:     &txnHash,
+		SpenderInputIndex: outpoint.Index,
+	}
+
+	// Cache details in case of late registration.
+	if _, ok := m.spends[*outpoint]; !ok {
+		m.spends[*outpoint] = details
+	}
+
+	// Deliver any backlogged spend notifications.
 	if spendChans, ok := m.spendMap[*outpoint]; ok {
 		delete(m.spendMap, *outpoint)
 		for _, spendChan := range spendChans {
-			txnHash := txn.TxHash()
 			spendChan <- &chainntnfs.SpendDetail{
-				SpentOutPoint:     outpoint,
-				SpendingHeight:    height,
-				SpendingTx:        txn,
-				SpenderTxHash:     &txnHash,
-				SpenderInputIndex: outpoint.Index,
+				SpentOutPoint:     details.SpentOutPoint,
+				SpendingHeight:    details.SpendingHeight,
+				SpendingTx:        details.SpendingTx,
+				SpenderTxHash:     details.SpenderTxHash,
+				SpenderInputIndex: details.SpenderInputIndex,
 			}
 		}
 	}