From a1cac6d54b2d80087aa53bd4377df30538879db5 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Thu, 12 Jan 2023 21:17:28 +0800
Subject: [PATCH] htlcswitch: use mock clock in `TestMailBoxAddExpiry`

This commit replaces the clock used in the unit test
`TestMailBoxAddExpiry`. Previously the `TestClock` is used, resulting in
the unit test being not so "unit" as the maintainer needs to know the
detailed implementation of `clock.Clock`, resulting in debugging a
failed unit test more difficult as the cognitive cost is high.
Re-implement `clock.Clock` also means we need to maintain more. This is
now solved by using mock clock so we can ignore the implementation
details and care only the returned results.
---
 docs/release-notes/release-notes-0.16.0.md |  3 +
 htlcswitch/mailbox_test.go                 | 75 +++++++++++++++++-----
 2 files changed, 61 insertions(+), 17 deletions(-)

diff --git a/docs/release-notes/release-notes-0.16.0.md b/docs/release-notes/release-notes-0.16.0.md
index 47f369ce6..3e5b81825 100644
--- a/docs/release-notes/release-notes-0.16.0.md
+++ b/docs/release-notes/release-notes-0.16.0.md
@@ -205,6 +205,9 @@ certain large transactions](https://github.com/lightningnetwork/lnd/pull/7100).
 * [Code style cleanup](https://github.com/lightningnetwork/lnd/pull/7308) in the
   funding package.
 
+* [Fixed a flake in the TestMailBoxAddExpiry unit
+  test](https://github.com/lightningnetwork/lnd/pull/7314).
+
 ## `lncli`
 
 * [Add an `insecure` flag to skip tls auth as well as a `metadata` string slice
diff --git a/htlcswitch/mailbox_test.go b/htlcswitch/mailbox_test.go
index 9acb84d5f..e48aacfcd 100644
--- a/htlcswitch/mailbox_test.go
+++ b/htlcswitch/mailbox_test.go
@@ -10,8 +10,10 @@ import (
 	"github.com/davecgh/go-spew/spew"
 	"github.com/lightningnetwork/lnd/channeldb"
 	"github.com/lightningnetwork/lnd/clock"
+	"github.com/lightningnetwork/lnd/lnmock"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -192,6 +194,35 @@ type mailboxContext struct {
 	forwards chan *htlcPacket
 }
 
+// newMailboxContextWithClock creates a new mailbox context with the given
+// mocked clock.
+//
+// TODO(yy): replace all usage of `newMailboxContext` with this method.
+func newMailboxContextWithClock(t *testing.T,
+	clock clock.Clock) *mailboxContext {
+
+	ctx := &mailboxContext{
+		t:        t,
+		forwards: make(chan *htlcPacket, 1),
+	}
+
+	failMailboxUpdate := func(outScid,
+		mboxScid lnwire.ShortChannelID) lnwire.FailureMessage {
+
+		return &lnwire.FailTemporaryNodeFailure{}
+	}
+
+	ctx.mailbox = newMemoryMailBox(&mailBoxConfig{
+		failMailboxUpdate: failMailboxUpdate,
+		forwardPackets:    ctx.forward,
+		clock:             clock,
+	})
+	ctx.mailbox.Start()
+	t.Cleanup(ctx.mailbox.Stop)
+
+	return ctx
+}
+
 func newMailboxContext(t *testing.T, startTime time.Time,
 	expiry time.Duration) *mailboxContext {
 
@@ -461,34 +492,44 @@ func TestMailBoxPacketPrioritization(t *testing.T) {
 	}
 }
 
-// TestMailBoxAddExpiry asserts that the mailbox will cancel back Adds that have
-// reached their expiry time.
+// TestMailBoxAddExpiry asserts that the mailbox will cancel back Adds that
+// have reached their expiry time.
 func TestMailBoxAddExpiry(t *testing.T) {
-	var (
-		expiry            = time.Minute
-		batchDelay        = time.Second
-		firstBatchStart   = time.Now()
-		firstBatchExpiry  = firstBatchStart.Add(expiry)
-		secondBatchStart  = firstBatchStart.Add(batchDelay)
-		secondBatchExpiry = secondBatchStart.Add(expiry)
-	)
-
-	ctx := newMailboxContext(t, firstBatchStart, expiry)
-
 	// Each batch will consist of 10 messages.
 	const numBatchPackets = 10
 
-	firstBatch := ctx.sendAdds(0, numBatchPackets)
+	// deadline is the returned value from the `pktWithExpiry.deadline`.
+	deadline := make(chan time.Time, numBatchPackets*2)
 
-	ctx.clock.SetTime(secondBatchStart)
+	// Create a mock clock and mock the methods.
+	mockClock := &lnmock.MockClock{}
+	mockClock.On("Now").Return(time.Now())
+
+	// Mock TickAfter, which mounts the above `deadline` channel to the
+	// returned value from `pktWithExpiry.deadline`.
+	mockClock.On("TickAfter", mock.Anything).Return(deadline)
+
+	// Create a test mailbox context.
+	ctx := newMailboxContextWithClock(t, mockClock)
+
+	// Send 10 packets and assert no failures are sent back.
+	firstBatch := ctx.sendAdds(0, numBatchPackets)
 	ctx.checkFails(nil)
 
+	// Send another 10 packets and assert no failures are sent back.
 	secondBatch := ctx.sendAdds(numBatchPackets, numBatchPackets)
+	ctx.checkFails(nil)
 
-	ctx.clock.SetTime(firstBatchExpiry)
+	// Tick 10 times and we should see the first batch expired.
+	for i := 0; i < numBatchPackets; i++ {
+		deadline <- time.Now()
+	}
 	ctx.checkFails(firstBatch)
 
-	ctx.clock.SetTime(secondBatchExpiry)
+	// Tick another 10 times and we should see the second batch expired.
+	for i := 0; i < numBatchPackets; i++ {
+		deadline <- time.Now()
+	}
 	ctx.checkFails(secondBatch)
 }