From a259317d72a528d4fd0671b9c1fef29374732f2b Mon Sep 17 00:00:00 2001
From: Elliott Jin <elliott.jin@gmail.com>
Date: Tue, 16 Feb 2021 09:55:22 -0800
Subject: [PATCH] netann: add test for RequestAuto

---
 netann/chan_status_manager_test.go | 52 ++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/netann/chan_status_manager_test.go b/netann/chan_status_manager_test.go
index e9ff61d6..0b79597d 100644
--- a/netann/chan_status_manager_test.go
+++ b/netann/chan_status_manager_test.go
@@ -427,6 +427,18 @@ func (h *testHarness) assertDisables(channels []*channeldb.OpenChannel, expErr e
 	}
 }
 
+// assertAutos requests auto state management for all of the passed channels, and
+// asserts that the errors returned from RequestAuto matches expErr.
+func (h *testHarness) assertAutos(channels []*channeldb.OpenChannel,
+	expErr error) {
+
+	h.t.Helper()
+
+	for _, channel := range channels {
+		h.assertAuto(channel.FundingOutpoint, expErr)
+	}
+}
+
 // assertEnable requests an enable for the given outpoint, and asserts that the
 // returned error matches expErr.
 func (h *testHarness) assertEnable(outpoint wire.OutPoint, expErr error,
@@ -453,6 +465,17 @@ func (h *testHarness) assertDisable(outpoint wire.OutPoint, expErr error,
 	}
 }
 
+// assertAuto requests auto state management for the given outpoint, and asserts
+// that the returned error matches expErr.
+func (h *testHarness) assertAuto(outpoint wire.OutPoint, expErr error) {
+	h.t.Helper()
+
+	err := h.mgr.RequestAuto(outpoint)
+	if err != expErr {
+		h.t.Fatalf("expected error: %v, got %v", expErr, err)
+	}
+}
+
 // assertNoUpdates waits for the specified duration, and asserts that no updates
 // are announced on the network.
 func (h *testHarness) assertNoUpdates(duration time.Duration) {
@@ -844,6 +867,35 @@ var stateMachineTests = []stateMachineTest{
 			// Request enables with manual = true should succeed.
 			h.assertEnables(h.graph.chans(), nil, true)
 
+			// Expect to see them all enabled on the network again.
+			h.assertUpdates(
+				h.graph.chans(), true, h.safeDisableTimeout,
+			)
+		},
+	},
+	{
+		name:         "restore auto",
+		startActive:  true,
+		startEnabled: true,
+		fn: func(h testHarness) {
+			// Request manual disables for all channels.
+			h.assertDisables(h.graph.chans(), nil, true)
+
+			// Expect to see them all disabled on the network.
+			h.assertUpdates(
+				h.graph.chans(), false, h.safeDisableTimeout,
+			)
+
+			// Request enables with manual = false should fail.
+			h.assertEnables(
+				h.graph.chans(), netann.ErrEnableManuallyDisabledChan, false,
+			)
+
+			// Request enables with manual = false should succeed after
+			// restoring auto state management.
+			h.assertAutos(h.graph.chans(), nil)
+			h.assertEnables(h.graph.chans(), nil, false)
+
 			// Expect to see them all enabled on the network again.
 			h.assertUpdates(
 				h.graph.chans(), true, h.safeDisableTimeout,