From 03bf786b47313f58a20c338a49cc6a1e068403c1 Mon Sep 17 00:00:00 2001
From: tbjump <unknown>
Date: Sat, 15 Jul 2023 23:05:51 +0000
Subject: [PATCH] node/db: Support in-memory DB and use in tests

---
 node/pkg/db/db.go          |  1 +
 node/pkg/db/open.go        | 22 +++++++++++++++++-----
 node/pkg/node/node_test.go | 11 ++---------
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/node/pkg/db/db.go b/node/pkg/db/db.go
index b038a326f..ce2a51ac1 100644
--- a/node/pkg/db/db.go
+++ b/node/pkg/db/db.go
@@ -83,6 +83,7 @@ func (i *VAAID) EmitterPrefixBytes() []byte {
 	return []byte(fmt.Sprintf("signed/%d/%s", i.EmitterChain, i.EmitterAddress))
 }
 
+// TODO: Deprecate in favor of OpenDb
 func Open(path string) (*Database, error) {
 	db, err := badger.Open(badger.DefaultOptions(path))
 	if err != nil {
diff --git a/node/pkg/db/open.go b/node/pkg/db/open.go
index 458924e22..90890627d 100644
--- a/node/pkg/db/open.go
+++ b/node/pkg/db/open.go
@@ -4,18 +4,30 @@ import (
 	"os"
 	"path"
 
+	"github.com/dgraph-io/badger/v3"
 	"go.uber.org/zap"
 )
 
 func OpenDb(logger *zap.Logger, dataDir *string) *Database {
-	dbPath := path.Join(*dataDir, "db")
-	if err := os.MkdirAll(dbPath, 0700); err != nil {
-		logger.Fatal("failed to create database directory", zap.Error(err))
+	var options badger.Options
+
+	if dataDir != nil {
+		dbPath := path.Join(*dataDir, "db")
+		if err := os.MkdirAll(dbPath, 0700); err != nil {
+			logger.Fatal("failed to create database directory", zap.Error(err))
+		}
+
+		options = badger.DefaultOptions(dbPath)
+	} else {
+		options = badger.DefaultOptions("").WithInMemory(true)
 	}
-	db, err := Open(dbPath)
+
+	db, err := badger.Open(options)
 	if err != nil {
 		logger.Fatal("failed to open database", zap.Error(err))
 	}
 
-	return db
+	return &Database{
+		db: db,
+	}
 }
diff --git a/node/pkg/node/node_test.go b/node/pkg/node/node_test.go
index 2ab5da371..ed5317624 100644
--- a/node/pkg/node/node_test.go
+++ b/node/pkg/node/node_test.go
@@ -13,7 +13,6 @@ import (
 	math_rand "math/rand"
 	"net/http"
 	"os"
-	"path"
 	"regexp"
 	"strconv"
 	"strings"
@@ -128,9 +127,7 @@ func mockGuardianRunnable(gs []*mockGuardian, mockGuardianIndex uint, obsDb mock
 		logger := supervisor.Logger(ctx)
 
 		// setup db
-		dataDir := fmt.Sprintf("/tmp/test_guardian_%d", mockGuardianIndex)
-		_ = os.RemoveAll(dataDir) // delete any pre-existing data
-		db := db.OpenDb(logger, &dataDir)
+		db := db.OpenDb(logger, nil)
 		defer db.Close()
 
 		// set environment
@@ -174,7 +171,7 @@ func mockGuardianRunnable(gs []*mockGuardian, mockGuardianIndex uint, obsDb mock
 			GuardianOptionP2P(gs[mockGuardianIndex].p2pKey, networkID, bootstrapPeers, nodeName, false, p2pPort, func() string { return "" }),
 			GuardianOptionPublicRpcSocket(publicSocketPath, common.GrpcLogDetailFull),
 			GuardianOptionPublicrpcTcpService(publicRpc, common.GrpcLogDetailFull),
-			GuardianOptionPublicWeb(mockPublicWeb(mockGuardianIndex), publicSocketPath, "", false, path.Join(dataDir, "autocert")),
+			GuardianOptionPublicWeb(mockPublicWeb(mockGuardianIndex), publicSocketPath, "", false, ""),
 			GuardianOptionAdminService(adminSocketPath, nil, nil, rpcMap),
 			GuardianOptionStatusServer(fmt.Sprintf("[::]:%d", mockStatusPort(mockGuardianIndex))),
 			GuardianOptionProcessor(),
@@ -190,10 +187,6 @@ func mockGuardianRunnable(gs []*mockGuardian, mockGuardianIndex uint, obsDb mock
 		}
 
 		<-ctx.Done()
-
-		// cleanup
-		// _ = os.RemoveAll(dataDir) // we don't do this for now since this could run before BadgerDB's flush(), causing an error; Meh
-
 		return nil
 	}
 }