Adjust unit tests to store.close()

Author: jasonjoo2010

core/heartbeat_test.go

@@ -32,6 +32,9 @@ func newManager(t *testing.T, store store.Store) *ScheduleManager {
 
 func TestHeartbeat(t *testing.T) {
 	store := memory.New()
+	defer func() {
+		assert.Nil(t, store.Close())
+	}()
 	manager1 := newManager(t, store)
 	manager2 := newManager(t, store)
 

core/leader_test.go

@@ -39,6 +39,9 @@ func (demo *DemoWorker) Stop(strategyId, parameter string) error {
 func TestStopAllWorkers(t *testing.T) {
 	// stop parellel
 	store := memory.New()
+	defer func() {
+		assert.Nil(t, store.Close())
+	}()
 	manager := newManager(t, store)
 	demo := &DemoWorker{}
 	worker.RegisterInstName("demoStopWorker", demo)
@@ -68,6 +71,9 @@ func TestStopAllWorkers(t *testing.T) {
 
 func TestLeader(t *testing.T) {
 	store := memory.New()
+	defer func() {
+		assert.Nil(t, store.Close())
+	}()
 	manager1 := newManager(t, store)
 	manager2 := newManager(t, store)
 
@@ -98,6 +104,9 @@ func TestLeader(t *testing.T) {
 
 func TestExpiredSchedulers(t *testing.T) {
 	store := memory.New()
+	defer func() {
+		assert.Nil(t, store.Close())
+	}()
 	manager := newManager(t, store)
 	managerExpired := newManager(t, store)
 
@@ -136,6 +145,9 @@ func TestExpiredSchedulers(t *testing.T) {
 
 func TestGenerateRuntimes(t *testing.T) {
 	store := memory.New()
+	defer func() {
+		assert.Nil(t, store.Close())
+	}()
 	manager := newManager(t, store)
 
 	store.CreateStrategy(&definition.Strategy{
@@ -164,6 +176,9 @@ func TestGenerateRuntimes(t *testing.T) {
 func TestAssign(t *testing.T) {
 	worker.RegisterName("demo", &DemoWorker{})
 	store := memory.New()
+	defer func() {
+		assert.Nil(t, store.Close())
+	}()
 	manager1 := newManager(t, store)
 	manager2 := newManager(t, store)
 	manager3 := newManager(t, store)

store/memory/memory.go

@@ -6,6 +6,7 @@ package memory
 
 import (
 	"encoding/json"
+	"errors"
 	"strconv"
 	"strings"
 	"sync"
@@ -20,6 +21,7 @@ import (
 type MemoryStore struct {
 	sequence        uint64
 	mutex           *sync.Mutex
+	closed          bool
 	taskItemsConfig map[string]int64
 	tasks           map[string]*definition.Task
 	strategies      map[string]*definition.Strategy
@@ -77,6 +79,13 @@ func (s *MemoryStore) Time() int64 {
 }
 
 func (s *MemoryStore) Close() error {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+	if s.closed {
+		return errors.New("The store has been already closed")
+	}
+
+	s.closed = true
 	return nil
 }