lntest: add more methods to assert mempool state

lntest/harness.go

@@ -1610,6 +1610,57 @@ func (h *HarnessTest) MineBlocksAndAssertNumTxes(num uint32,
 	return blocks
 }
 
+// cleanMempool mines blocks till the mempool is empty and asserts all active
+// nodes have synced to the chain.
+func (h *HarnessTest) cleanMempool() {
+	_, startHeight := h.Miner.GetBestBlock()
+
+	// Mining the blocks slow to give `lnd` more time to sync.
+	var bestBlock *wire.MsgBlock
+	err := wait.NoError(func() error {
+		// If mempool is empty, exit.
+		mem := h.Miner.GetRawMempool()
+		if len(mem) == 0 {
+			_, height := h.Miner.GetBestBlock()
+			h.Logf("Mined %d blocks when cleanup the mempool",
+				height-startHeight)
+
+			return nil
+		}
+
+		// Otherwise mine a block.
+		blocks := h.Miner.MineBlocksSlow(1)
+		bestBlock = blocks[len(blocks)-1]
+
+		return fmt.Errorf("still have %d txes in mempool", len(mem))
+	}, wait.MinerMempoolTimeout)
+	require.NoError(h, err, "timeout cleaning up mempool")
+
+	// Exit early if the best block is nil, which means we haven't mined
+	// any blocks during the cleanup.
+	if bestBlock == nil {
+		return
+	}
+
+	// Make sure all the active nodes are synced.
+	h.AssertActiveNodesSyncedTo(bestBlock)
+}
+
+// CleanShutDown is used to quickly end a test by shutting down all non-standby
+// nodes and mining blocks to empty the mempool.
+//
+// NOTE: this method provides a faster exit for a test that involves force
+// closures as the caller doesn't need to mine all the blocks to make sure the
+// mempool is empty.
+func (h *HarnessTest) CleanShutDown() {
+	// First, shutdown all non-standby nodes to prevent new transactions
+	// being created and fed into the mempool.
+	h.shutdownNonStandbyNodes()
+
+	// Now mine blocks till the mempool is empty.
+	h.cleanMempool()
+}
+
 // MineEmptyBlocks mines a given number of empty blocks.
 //
 // NOTE: this differs from miner's `MineEmptyBlocks` as it requires the nodes

lntest/harness_assertion.go

@@ -887,6 +887,15 @@ func (h *HarnessTest) Random32Bytes() []byte {
 	return randBuf
 }
 
+// RandomPreimage generates a random preimage which can be used as a payment
+// preimage.
+func (h *HarnessTest) RandomPreimage() lntypes.Preimage {
+	var preimage lntypes.Preimage
+	copy(preimage[:], h.Random32Bytes())
+
+	return preimage
+}
+
 // DecodeAddress decodes a given address and asserts there's no error.
 func (h *HarnessTest) DecodeAddress(addr string) btcutil.Address {
 	resp, err := btcutil.DecodeAddress(addr, harnessNetParams)
@@ -1256,6 +1265,111 @@ func (h *HarnessTest) AssertActiveHtlcs(hn *node.HarnessNode,
 	require.NoError(h, err, "timeout checking active HTLCs")
 }
 
+// AssertIncomingHTLCActive asserts the node has a pending incoming HTLC in the
+// given channel. Returns the HTLC if found and active.
+func (h *HarnessTest) AssertIncomingHTLCActive(hn *node.HarnessNode,
+	cp *lnrpc.ChannelPoint, payHash []byte) *lnrpc.HTLC {
+
+	return h.assertHLTCActive(hn, cp, payHash, true)
+}
+
+// AssertOutgoingHTLCActive asserts the node has a pending outgoing HTLC in the
+// given channel. Returns the HTLC if found and active.
+func (h *HarnessTest) AssertOutgoingHTLCActive(hn *node.HarnessNode,
+	cp *lnrpc.ChannelPoint, payHash []byte) *lnrpc.HTLC {
+
+	return h.assertHLTCActive(hn, cp, payHash, false)
+}
+
+// assertHLTCActive asserts the node has a pending HTLC in the given channel.
+// Returns the HTLC if found and active.
+func (h *HarnessTest) assertHLTCActive(hn *node.HarnessNode,
+	cp *lnrpc.ChannelPoint, payHash []byte, incoming bool) *lnrpc.HTLC {
+
+	var result *lnrpc.HTLC
+	target := hex.EncodeToString(payHash)
+
+	err := wait.NoError(func() error {
+		// We require the RPC call to be succeeded and won't wait for
+		// it as it's an unexpected behavior.
+		ch := h.GetChannelByChanPoint(hn, cp)
+
+		// Check all payment hashes active for this channel.
+		for _, htlc := range ch.PendingHtlcs {
+			h := hex.EncodeToString(htlc.HashLock)
+			if h != target {
+				continue
+			}
+
+			// If the payment hash is found, check the incoming
+			// field.
+			if htlc.Incoming == incoming {
+				// Found it and return.
+				result = htlc
+				return nil
+			}
+
+			// Otherwise we do have the HTLC but its direction is
+			// not right.
+			have, want := "outgoing", "incoming"
+			if htlc.Incoming {
+				have, want = "incoming", "outgoing"
+			}
+
+			return fmt.Errorf("node[%s] have htlc(%v), want: %s, "+
+				"have: %s", hn.Name(), payHash, want, have)
+		}
+
+		return fmt.Errorf("node [%s:%x] didn't have: the payHash %v",
+			hn.Name(), hn.PubKey[:], payHash)
+	}, DefaultTimeout)
+	require.NoError(h, err, "timeout checking pending HTLC")
+
+	return result
+}
+
+// AssertHLTCNotActive asserts the node doesn't have a pending HTLC in the
+// given channel, which mean either the HTLC never exists, or it was pending
+// and now settled. Returns the HTLC if found and active.
+//
+// NOTE: to check a pending HTLC becoming settled, first use AssertHLTCActive
+// then follow this check.
+func (h *HarnessTest) AssertHLTCNotActive(hn *node.HarnessNode,
+	cp *lnrpc.ChannelPoint, payHash []byte) *lnrpc.HTLC {
+
+	var result *lnrpc.HTLC
+	target := hex.EncodeToString(payHash)
+
+	err := wait.NoError(func() error {
+		// We require the RPC call to be succeeded and won't wait for
+		// it as it's an unexpected behavior.
+		ch := h.GetChannelByChanPoint(hn, cp)
+
+		// Check all payment hashes active for this channel.
+		for _, htlc := range ch.PendingHtlcs {
+			h := hex.EncodeToString(htlc.HashLock)
+
+			// Break if found the htlc.
+			if h == target {
+				result = htlc
+				break
+			}
+		}
+
+		// If we've found nothing, we're done.
+		if result == nil {
+			return nil
+		}
+
+		// Otherwise return an error.
+		return fmt.Errorf("node [%s:%x] still has: the payHash %x",
+			hn.Name(), hn.PubKey[:], payHash)
+	}, DefaultTimeout)
+	require.NoError(h, err, "timeout checking pending HTLC")
+
+	return result
+}
+
 // ReceiveSingleInvoice waits until a message is received on the subscribe
 // single invoice stream or the timeout is reached.
 func (h *HarnessTest) ReceiveSingleInvoice(
@@ -1450,16 +1564,17 @@ func (h *HarnessTest) AssertPaymentStatus(hn *node.HarnessNode,
 	status lnrpc.Payment_PaymentStatus) *lnrpc.Payment {
 
 	var target *lnrpc.Payment
+	payHash := preimage.Hash()
 
 	err := wait.NoError(func() error {
-		p := h.findPayment(hn, preimage.Hash().String())
+		p := h.findPayment(hn, payHash.String())
 		if status == p.Status {
 			target = p
 			return nil
 		}
 
 		return fmt.Errorf("payment: %v status not match, want %s "+
-			"got %s", preimage, status, p.Status)
+			"got %s", payHash, status, p.Status)
 	}, DefaultTimeout)
 	require.NoError(h, err, "timeout checking payment status")
 
@@ -2349,3 +2464,53 @@ func (h *HarnessTest) AssertWalletAccountBalance(hn *node.HarnessNode,
 	}, DefaultTimeout)
 	require.NoError(h, err, "timeout checking wallet account balance")
 }
+
+// AssertClosingTxInMempool assert that the closing transaction of the given
+// channel point can be found in the mempool. If the channel has anchors, it
+// will assert the anchor sweep tx is also in the mempool.
+func (h *HarnessTest) AssertClosingTxInMempool(cp *lnrpc.ChannelPoint,
+	c lnrpc.CommitmentType) *wire.MsgTx {
+
+	// Get expected number of txes to be found in the mempool.
+	expectedTxes := 1
+	hasAnchors := CommitTypeHasAnchors(c)
+	if hasAnchors {
+		expectedTxes = 2
+	}
+
+	// Wait for the expected txes to be found in the mempool.
+	h.Miner.AssertNumTxsInMempool(expectedTxes)
+
+	// Get the closing tx from the mempool.
+	op := h.OutPointFromChannelPoint(cp)
+	closeTx := h.Miner.AssertOutpointInMempool(op)
+
+	return closeTx
+}
+
+// AssertClosingTxInMempool assert that the closing transaction of the given
+// channel point can be found in the mempool. If the channel has anchors, it
+// will assert the anchor sweep tx is also in the mempool.
+func (h *HarnessTest) MineClosingTx(cp *lnrpc.ChannelPoint,
+	c lnrpc.CommitmentType) *wire.MsgTx {
+
+	// Get expected number of txes to be found in the mempool.
+	expectedTxes := 1
+	hasAnchors := CommitTypeHasAnchors(c)
+	if hasAnchors {
+		expectedTxes = 2
+	}
+
+	// Wait for the expected txes to be found in the mempool.
+	h.Miner.AssertNumTxsInMempool(expectedTxes)
+
+	// Get the closing tx from the mempool.
+	op := h.OutPointFromChannelPoint(cp)
+	closeTx := h.Miner.AssertOutpointInMempool(op)
+
+	// Mine a block to confirm the closing transaction and potential anchor
+	// sweep.
+	h.MineBlocksAndAssertNumTxes(1, expectedTxes)
+
+	return closeTx
+}

lntest/harness_miner.go

@@ -301,6 +301,40 @@ func (h *HarnessMiner) AssertTxInMempool(txid *chainhash.Hash) *wire.MsgTx {
 	return msgTx
 }
 
+// AssertTxNotInMempool asserts a given transaction cannot be found in the
+// mempool. It assumes the mempool is not empty.
+//
+// NOTE: this should be used after `AssertTxInMempool` to ensure the tx has
+// entered the mempool before. Otherwise it might give false positive and the
+// tx may enter the mempool after the check.
+func (h *HarnessMiner) AssertTxNotInMempool(txid chainhash.Hash) *wire.MsgTx {
+	var msgTx *wire.MsgTx
+
+	err := wait.NoError(func() error {
+		// We require the RPC call to be succeeded and won't wait for
+		// it as it's an unexpected behavior.
+		mempool := h.GetRawMempool()
+
+		if len(mempool) == 0 {
+			return fmt.Errorf("empty mempool")
+		}
+
+		for _, memTx := range mempool {
+			// Check the values are equal.
+			if txid.IsEqual(memTx) {
+				return fmt.Errorf("expect txid %v to be NOT "+
+					"found in mempool", txid)
+			}
+		}
+
+		return nil
+	}, wait.MinerMempoolTimeout)
+
+	require.NoError(h, err, "timeout checking tx not in mempool")
+
+	return msgTx
+}
+
 // SendOutputsWithoutChange uses the miner to send the given outputs using the
 // specified fee rate and returns the txid.
 func (h *HarnessMiner) SendOutputsWithoutChange(outputs []*wire.TxOut,