From 63b0ac8b354a1b33305b950c596e31dbafb2fd3e Mon Sep 17 00:00:00 2001
From: TwinProduction <twin@twinnation.org>
Date: Wed, 19 May 2021 00:55:03 -0400
Subject: [PATCH] Improve test coverage

---
 storage/storage.go      |  8 ++++----
 storage/storage_test.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 4 deletions(-)
 create mode 100644 storage/storage_test.go

diff --git a/storage/storage.go b/storage/storage.go
index 9adc1f87..09cf0061 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -37,6 +37,10 @@ func Get() store.Store {
 func Initialize(cfg *Config) error {
 	initialized = true
 	var err error
+	if cancelFunc != nil {
+		// Stop the active autoSave task
+		cancelFunc()
+	}
 	if cfg == nil || len(cfg.File) == 0 {
 		log.Println("[storage][Initialize] Creating storage provider")
 		provider, err = memory.NewStore("")
@@ -44,10 +48,6 @@ func Initialize(cfg *Config) error {
 			return err
 		}
 	} else {
-		if cancelFunc != nil {
-			// Stop the active autoSave task
-			cancelFunc()
-		}
 		ctx, cancelFunc = context.WithCancel(context.Background())
 		log.Printf("[storage][Initialize] Creating storage provider with file=%s", cfg.File)
 		provider, err = memory.NewStore(cfg.File)
diff --git a/storage/storage_test.go b/storage/storage_test.go
new file mode 100644
index 00000000..92359cb7
--- /dev/null
+++ b/storage/storage_test.go
@@ -0,0 +1,37 @@
+package storage
+
+import (
+	"testing"
+	"time"
+)
+
+func TestInitialize(t *testing.T) {
+	file := t.TempDir() + "/test.db"
+	err := Initialize(&Config{File: file})
+	if err != nil {
+		t.Fatal("shouldn't have returned an error")
+	}
+	if cancelFunc == nil {
+		t.Error("cancelFunc shouldn't have been nil")
+	}
+	if ctx == nil {
+		t.Error("ctx shouldn't have been nil")
+	}
+	// Try to initialize it again
+	err = Initialize(&Config{File: file})
+	if err != nil {
+		t.Fatal("shouldn't have returned an error")
+	}
+	cancelFunc()
+}
+
+func TestAutoSave(t *testing.T) {
+	file := t.TempDir() + "/test.db"
+	if err := Initialize(&Config{File: file}); err != nil {
+		t.Fatal("shouldn't have returned an error")
+	}
+	go autoSave(3*time.Millisecond, ctx)
+	time.Sleep(15 * time.Millisecond)
+	cancelFunc()
+	time.Sleep(5 * time.Millisecond)
+}