From a12c0c4bd7c51cb9348da309e36695a61cda73b4 Mon Sep 17 00:00:00 2001
From: Benjamin Wang <wachao@vmware.com>
Date: Sat, 25 Mar 2023 13:56:57 +0800
Subject: [PATCH] test: support enabling strict mode in testing

Signed-off-by: Benjamin Wang <wachao@vmware.com>
---
 Makefile                        | 10 +++++++++-
 internal/btesting/btesting.go   | 20 ++++++++++++++++++--
 internal/tests/tx_check_test.go |  4 ++++
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/Makefile b/Makefile
index e2a594e..e234574 100644
--- a/Makefile
+++ b/Makefile
@@ -18,6 +18,14 @@ ifdef TIMEOUT
 	TESTFLAGS_TIMEOUT=$(TIMEOUT)
 endif
 
+TESTFLAGS_ENABLE_STRICT_MODE=false
+ifdef ENABLE_STRICT_MODE
+	TESTFLAGS_ENABLE_STRICT_MODE=$(ENABLE_STRICT_MODE)
+endif
+
+.EXPORT_ALL_VARIABLES:
+TEST_ENABLE_STRICT_MODE=${TESTFLAGS_ENABLE_STRICT_MODE}
+
 .PHONY: fmt
 fmt:
 	!(gofmt -l -s -d $(shell find . -name \*.go) | grep '[a-z]')
@@ -53,7 +61,7 @@ gofail-enable: install-gofail
 	gofail enable .
 
 .PHONY: gofail-disable
-gofail-disable:
+gofail-disable: install-gofail
 	gofail disable .
 
 .PHONY: install-gofail
diff --git a/internal/btesting/btesting.go b/internal/btesting/btesting.go
index e9ef64b..ffa0d8b 100644
--- a/internal/btesting/btesting.go
+++ b/internal/btesting/btesting.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"strings"
 	"testing"
 	"time"
 
@@ -17,8 +18,12 @@ import (
 
 var statsFlag = flag.Bool("stats", false, "show performance stats")
 
-// TestFreelistType is used as a env variable for test to indicate the backend type
-const TestFreelistType = "TEST_FREELIST_TYPE"
+const (
+	// TestFreelistType is used as an env variable for test to indicate the backend type.
+	TestFreelistType = "TEST_FREELIST_TYPE"
+	// TestEnableStrictMode is used to enable strict check by default after opening each DB.
+	TestEnableStrictMode = "TEST_ENABLE_STRICT_MODE"
+)
 
 // DB is a test wrapper for bolt.DB.
 type DB struct {
@@ -60,6 +65,7 @@ func MustOpenDBWithOption(t testing.TB, f string, o *bolt.Options) *DB {
 		o:  o,
 		t:  t,
 	}
+	resDB.strictModeEnabledDefault()
 	t.Cleanup(resDB.PostTestCleanup)
 	return resDB
 }
@@ -113,6 +119,7 @@ func (db *DB) MustReopen() {
 	indb, err := bolt.Open(db.Path(), 0666, db.o)
 	require.NoError(db.t, err)
 	db.DB = indb
+	db.strictModeEnabledDefault()
 }
 
 // MustCheck runs a consistency check on the database and panics if any errors are found.
@@ -204,3 +211,12 @@ func (db *DB) PrintStats() {
 func truncDuration(d time.Duration) string {
 	return regexp.MustCompile(`^(\d+)(\.\d+)`).ReplaceAllString(d.String(), "$1")
 }
+
+func (db *DB) strictModeEnabledDefault() {
+	strictModeEnabled := strings.ToLower(os.Getenv(TestEnableStrictMode))
+	db.StrictMode = strictModeEnabled == "true"
+}
+
+func (db *DB) ForceDisableStrictMode() {
+	db.StrictMode = false
+}
diff --git a/internal/tests/tx_check_test.go b/internal/tests/tx_check_test.go
index 0476d01..7a4ab86 100644
--- a/internal/tests/tx_check_test.go
+++ b/internal/tests/tx_check_test.go
@@ -14,6 +14,7 @@ import (
 
 func TestTx_RecursivelyCheckPages_MisplacedPage(t *testing.T) {
 	db := btesting.MustCreateDB(t)
+	db.ForceDisableStrictMode()
 	require.NoError(t,
 		db.Fill([]byte("data"), 1, 10000,
 			func(tx int, k int) []byte { return []byte(fmt.Sprintf("%04d", k)) },
@@ -36,6 +37,7 @@ func TestTx_RecursivelyCheckPages_MisplacedPage(t *testing.T) {
 	require.NoError(t, surgeon.CopyPage(db.Path(), srcPage, targetPage))
 
 	db.MustReopen()
+	db.ForceDisableStrictMode()
 	require.NoError(t, db.Update(func(tx *bolt.Tx) error {
 		// Collect all the errors.
 		var errors []error
@@ -51,6 +53,7 @@ func TestTx_RecursivelyCheckPages_MisplacedPage(t *testing.T) {
 
 func TestTx_RecursivelyCheckPages_CorruptedLeaf(t *testing.T) {
 	db := btesting.MustCreateDB(t)
+	db.ForceDisableStrictMode()
 	require.NoError(t,
 		db.Fill([]byte("data"), 1, 10000,
 			func(tx int, k int) []byte { return []byte(fmt.Sprintf("%04d", k)) },
@@ -72,6 +75,7 @@ func TestTx_RecursivelyCheckPages_CorruptedLeaf(t *testing.T) {
 	require.NoError(t, guts_cli.WritePage(db.Path(), pbuf))
 
 	db.MustReopen()
+	db.ForceDisableStrictMode()
 	require.NoError(t, db.Update(func(tx *bolt.Tx) error {
 		// Collect all the errors.
 		var errors []error