mock: support late registration of spend ntfns

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)
+	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,17 +172,31 @@ func (m *mockSpendNotifier) Spend(outpoint *wire.OutPoint, height int32,
 	m.mtx.Lock()
 	defer m.mtx.Unlock()
 
-	if spendChans, ok := m.spendMap[*outpoint]; ok {
-		delete(m.spendMap, *outpoint)
-		for _, spendChan := range spendChans {
 	txnHash := txn.TxHash()
-			spendChan <- &chainntnfs.SpendDetail{
+	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 {
+			spendChan <- &chainntnfs.SpendDetail{
+				SpentOutPoint:     details.SpentOutPoint,
+				SpendingHeight:    details.SpendingHeight,
+				SpendingTx:        details.SpendingTx,
+				SpenderTxHash:     details.SpenderTxHash,
+				SpenderInputIndex: details.SpenderInputIndex,
+			}
 		}
 	}
 }