From b164c603a9af3f123f32bf94a6138521518abdc6 Mon Sep 17 00:00:00 2001
From: Matheus Degiovani <opensource@matheusd.com>
Date: Wed, 18 Oct 2023 12:55:56 -0300
Subject: [PATCH] dcrdtest: Add option for keeping node dir

This adds a flag to the Harness to instruct it to keep the node dir
after TearDown and, unless that flag is set, removes the node dir.

This improves the reliability of the test harness by not leaving behind
extraneous, unneeded data by default.

When calling the TearDownInTest variant, if the test has already failed
then the node dir is kept to ease debugging test failures.
---
 dcrdtest/rpc_harness.go      | 26 +++++++++++++++++++++++
 dcrdtest/rpc_harness_test.go | 40 ++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/dcrdtest/rpc_harness.go b/dcrdtest/rpc_harness.go
index 36f0401..6194df0 100644
--- a/dcrdtest/rpc_harness.go
+++ b/dcrdtest/rpc_harness.go
@@ -67,6 +67,8 @@ type Harness struct {
 	maxConnRetries int
 	nodeNum        int
 
+	keepNodeDir bool
+
 	sync.Mutex
 }
 
@@ -214,6 +216,15 @@ func New(t *testing.T, activeNet *chaincfg.Params, handlers *rpcclient.Notificat
 	return h, nil
 }
 
+// SetKeepNodeDir sets the flag  in the Harness on whether to keep or remove
+// its node dir after TearDown is called.
+//
+// This is NOT safe for concurrent access and MUST be called from the same
+// goroutine that calls SetUp and TearDown.
+func (h *Harness) SetKeepNodeDir(keep bool) {
+	h.keepNodeDir = keep
+}
+
 // SetUp initializes the rpc test state. Initialization includes: starting up a
 // simnet node, creating a websockets client and connecting to the started
 // node, and finally: optionally generating and submitting a testchain with a
@@ -319,13 +330,28 @@ func (h *Harness) TearDown() error {
 		h.wallet = nil
 	}
 
+	if !h.keepNodeDir {
+		if err := os.RemoveAll(h.testNodeDir); err != nil {
+			log.Warnf("Unable to remove test node dir %s: %v", h.testNodeDir, err)
+		} else {
+			log.Debugf("Removed test node dir %s", h.testNodeDir)
+		}
+	}
+
 	return nil
 }
 
 // TearDownInTest performs the TearDown during a test, logging the error to the
 // test object. If the test has not yet failed and the TearDown itself fails,
 // then this fails the test.
+//
+// If the test has already failed, then the dir for node data is kept for manual
+// debugging.
 func (h *Harness) TearDownInTest(t testing.TB) {
+	if t.Failed() {
+		h.SetKeepNodeDir(true)
+	}
+
 	err := h.TearDown()
 	if err != nil {
 		errMsg := fmt.Sprintf("Unable to teardown dcrdtest harness: %v", err)
diff --git a/dcrdtest/rpc_harness_test.go b/dcrdtest/rpc_harness_test.go
index b29b5cc..687c176 100644
--- a/dcrdtest/rpc_harness_test.go
+++ b/dcrdtest/rpc_harness_test.go
@@ -703,6 +703,12 @@ func TestSetupTeardown(t *testing.T) {
 		t.Fatalf("Unexpected nb of active goroutines: got %d, want %d",
 			gotCount, wantCount)
 	}
+
+	// Node dir should've been removed.
+	if _, err := os.Stat(mainHarness.testNodeDir); !errors.Is(err, os.ErrNotExist) {
+		t.Fatalf("Unexpected Stat(testNodeDir) error: got %v, want %v",
+			err, os.ErrNotExist)
+	}
 }
 
 // TestSetupWithError tests that when the setup of an rpc harness fails, it
@@ -794,3 +800,37 @@ func TestSetupWithWrongDcrd(t *testing.T) {
 		t.Fatalf("Unexpected error during TearDown: %v", err)
 	}
 }
+
+// TestKeepNodeDir asserts that when the harness is set to keep the node dir,
+// it is not removed after TearDown().
+func TestKeepNodeDir(t *testing.T) {
+	// Add logging to ease debugging this test.
+	lw := loggerWriter{l: t}
+	bknd := slog.NewBackend(lw)
+	UseLogger(bknd.Logger("TEST"))
+	log.SetLevel(slog.LevelDebug)
+	defer UseLogger(slog.Disabled)
+
+	params := chaincfg.RegNetParams()
+	mainHarness, err := New(t, params, nil, nil)
+	if err != nil {
+		t.Fatalf("unable to create main harness: %v", err)
+	}
+
+	// Perform the setup. This should succeed.
+	ctx := testctx.WithTimeout(t, time.Second*30)
+	if err := mainHarness.SetUp(ctx, true, 2); err != nil {
+		t.Fatalf("Unexpected error in Setup(): got %v, want %v", err, nil)
+	}
+
+	// Set to keep the node dir.
+	mainHarness.SetKeepNodeDir(true)
+
+	// Perform the TearDown. This should not error and the node dir should
+	// be kept.
+	mainHarness.TearDownInTest(t)
+	if _, err := os.Stat(mainHarness.testNodeDir); err != nil {
+		t.Fatalf("Unexpected Stat(testNodeDir) error: got %v, want %v",
+			err, nil)
+	}
+}