From 9fc197d8b191de51cb4afc42bf8ffbb2217469b5 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Tue, 28 Jan 2020 06:43:44 -0800
Subject: [PATCH] routing/route: fix TestMPPHop comment

---
 routing/route/route_test.go | 41 +++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/routing/route/route_test.go b/routing/route/route_test.go
index df1b9fc3..2095430b 100644
--- a/routing/route/route_test.go
+++ b/routing/route/route_test.go
@@ -101,6 +101,47 @@ func TestMPPHop(t *testing.T) {
 	}
 }
 
+// TestAMPHop asserts that a Hop will encode a non-nil AMP to final nodes of an
+// MPP record is also present, and fail otherwise.
+func TestAMPHop(t *testing.T) {
+	t.Parallel()
+
+	hop := Hop{
+		ChannelID:        1,
+		OutgoingTimeLock: 44,
+		AmtToForward:     testAmt,
+		LegacyPayload:    false,
+		AMP:              record.NewAMP([32]byte{}, [32]byte{}, 3),
+	}
+
+	// Encoding an AMP record to an intermediate hop w/o an MPP record
+	// should result in a failure.
+	var b bytes.Buffer
+	err := hop.PackHopPayload(&b, 2)
+	if err != ErrAMPMissingMPP {
+		t.Fatalf("expected err: %v, got: %v",
+			ErrAMPMissingMPP, err)
+	}
+
+	// Encoding an AMP record to a final hop w/o an MPP record should result
+	// in a failure.
+	b.Reset()
+	err = hop.PackHopPayload(&b, 0)
+	if err != ErrAMPMissingMPP {
+		t.Fatalf("expected err: %v, got: %v",
+			ErrAMPMissingMPP, err)
+	}
+
+	// Encoding an AMP record to a final hop w/ an MPP record should be
+	// successful.
+	hop.MPP = record.NewMPP(testAmt, testAddr)
+	b.Reset()
+	err = hop.PackHopPayload(&b, 0)
+	if err != nil {
+		t.Fatalf("expected err: %v, got: %v", nil, err)
+	}
+}
+
 // TestPayloadSize tests the payload size calculation that is provided by Hop
 // structs.
 func TestPayloadSize(t *testing.T) {