From af076d8ff497e5a9ef44af259a37721be4537fb2 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Fri, 23 Dec 2022 10:10:27 +0200
Subject: [PATCH] channeldb: add shouldFail param to ApplyMigrationWithDB

Add a shouldFail boolean parameter to the migtest ApplyMigrationWithDB
in order to make it easier to test migration failures.
---
 channeldb/migration30/migration_test.go | 2 ++
 channeldb/migtest/migtest.go            | 9 ++++++---
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/channeldb/migration30/migration_test.go b/channeldb/migration30/migration_test.go
index c748ebac0..5ec91fdff 100644
--- a/channeldb/migration30/migration_test.go
+++ b/channeldb/migration30/migration_test.go
@@ -108,6 +108,7 @@ func TestMigrateRevocationLog(t *testing.T) {
 				beforeMigration,
 				afterMigration,
 				MigrateRevocationLog,
+				false,
 			)
 		})
 		if !success {
@@ -569,5 +570,6 @@ func BenchmarkMigration(b *testing.B) {
 
 			return MigrateRevocationLog(db)
 		},
+		false,
 	)
 }
diff --git a/channeldb/migtest/migtest.go b/channeldb/migtest/migtest.go
index 65e36308f..dacca7dac 100644
--- a/channeldb/migtest/migtest.go
+++ b/channeldb/migtest/migtest.go
@@ -90,7 +90,7 @@ func ApplyMigration(t *testing.T,
 // supplied migration functions to take a db instance and construct their own
 // database transactions.
 func ApplyMigrationWithDb(t testing.TB, beforeMigration, afterMigration,
-	migrationFunc func(db kvdb.Backend) error) {
+	migrationFunc func(db kvdb.Backend) error, shouldFail bool) {
 
 	t.Helper()
 
@@ -106,8 +106,11 @@ func ApplyMigrationWithDb(t testing.TB, beforeMigration, afterMigration,
 	}
 
 	// Apply migration.
-	if err := migrationFunc(cdb); err != nil {
-		t.Fatalf("migrationFunc error: %v", err)
+	err = migrationFunc(cdb)
+	if shouldFail {
+		require.Error(t, err)
+	} else {
+		require.NoError(t, err)
 	}
 
 	// If there's no afterMigration, exit here.