Fix test running in memory only

2 files changed, 25 insertions, 20 deletions

diff --git a/activitypub/inbox/worker_test.go b/activitypub/inbox/worker_test.go
index 95ec2fdaa..93b090eb2 100644
--- a/activitypub/inbox/worker_test.go
+++ b/activitypub/inbox/worker_test.go
@@ -1,9 +1,7 @@
 package inbox
 
 import (
-	"io/ioutil"
 	"net/url"
-	"os"
 	"testing"
 
 	"github.com/go-fed/activity/streams"
@@ -50,12 +48,7 @@ func makeFakePerson() vocab.ActivityStreamsPerson {
 }
 
 func TestMain(m *testing.M) {
-	dbFile, err := ioutil.TempFile(os.TempDir(), ":memory:")
-	if err != nil {
-		panic(err)
-	}
-
-	data.SetupPersistence(dbFile.Name())
+	data.SetupPersistence(":memory:")
 	data.SetServerURL("https://my.cool.site.biz")
 	persistence.Setup(data.GetDatastore())
 	m.Run()
diff --git a/core/data/data.go b/core/data/data.go
index c28e7dccf..f233d542c 100644
--- a/core/data/data.go
+++ b/core/data/data.go
@@ -37,18 +37,34 @@ func GetStore() *Datastore {
 
 // SetupPersistence will open the datastore and make it available.
 func SetupPersistence(file string) error {
-	// Create empty DB file if it doesn't exist.
-	if !utils.DoesFileExists(file) {
-		log.Traceln("Creating new database at", file)
+	// Allow support for in-memory databases for tests.
 
-		_, err := os.Create(file)
+	var db *sql.DB
+
+	if file == ":memory:" {
+		inMemoryDb, err := sql.Open("sqlite3", file)
 		if err != nil {
 			log.Fatal(err.Error())
 		}
-	}
+		db = inMemoryDb
+	} else {
+		// Create empty DB file if it doesn't exist.
+		if !utils.DoesFileExists(file) {
+			log.Traceln("Creating new database at", file)
+
+			_, err := os.Create(file)
+			if err != nil {
+				log.Fatal(err.Error())
+			}
+		}
 
-	db, err := sql.Open("sqlite3", fmt.Sprintf("file:%s?_cache_size=10000&cache=shared&_journal_mode=WAL", file))
-	db.SetMaxOpenConns(1)
+		onDiskDb, err := sql.Open("sqlite3", fmt.Sprintf("file:%s?_cache_size=10000&cache=shared&_journal_mode=WAL", file))
+		if err != nil {
+			return err
+		}
+		db = onDiskDb
+		db.SetMaxOpenConns(1)
+	}
 	_db = db
 
 	// Some SQLite optimizations
@@ -60,10 +76,6 @@ func SetupPersistence(file string) error {
 	createWebhooksTable()
 	createUsersTable(db)
 
-	if err != nil {
-		return err
-	}
-
 	if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS config (
 		"key" string NOT NULL PRIMARY KEY,
 		"value" TEXT
@@ -72,7 +84,7 @@ func SetupPersistence(file string) error {
 	}
 
 	var version int
-	err = db.QueryRow("SELECT value FROM config WHERE key='version'").
+	err := db.QueryRow("SELECT value FROM config WHERE key='version'").
 		Scan(&version)
 	if err != nil {
 		if err != sql.ErrNoRows {