bbgo: refactor isolation and add more tests

pkg/bbgo/graceful_shutdown.go

@@ -30,14 +30,14 @@ func (g *GracefulShutdown) Shutdown(ctx context.Context) {
 }
 
 func OnShutdown(ctx context.Context, f ShutdownHandler) {
-	isolatedContext := NewIsolationFromContext(ctx)
+	isolatedContext := GetIsolationFromContext(ctx)
 	isolatedContext.gracefulShutdown.OnShutdown(f)
 }
 
 func Shutdown(ctx context.Context) {
 	logrus.Infof("shutting down...")
 
-	isolatedContext := NewIsolationFromContext(ctx)
+	isolatedContext := GetIsolationFromContext(ctx)
 	todo := context.WithValue(context.TODO(), IsolationContextKey, isolatedContext)
 
 	timeoutCtx, cancel := context.WithTimeout(todo, 30*time.Second)

pkg/bbgo/isolation.go

@@ -8,21 +8,21 @@ import (
 
 const IsolationContextKey = "bbgo"
 
-var defaultIsolation = NewIsolation()
+var defaultIsolation = NewDefaultIsolation()
 
 type Isolation struct {
 	gracefulShutdown         GracefulShutdown
 	persistenceServiceFacade *service.PersistenceServiceFacade
 }
 
-func NewIsolation() *Isolation {
+func NewDefaultIsolation() *Isolation {
 	return &Isolation{
 		gracefulShutdown:         GracefulShutdown{},
-		persistenceServiceFacade: DefaultPersistenceServiceFacade,
+		persistenceServiceFacade: persistenceServiceFacade,
 	}
 }
 
-func NewIsolationFromContext(ctx context.Context) *Isolation {
+func GetIsolationFromContext(ctx context.Context) *Isolation {
 	isolatedContext, ok := ctx.Value(IsolationContextKey).(*Isolation)
 	if ok {
 		return isolatedContext

pkg/bbgo/isolation_test.go

@@ -0,0 +1,24 @@
+package bbgo
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetIsolationFromContext(t *testing.T) {
+	ctx := context.Background()
+	isolation := GetIsolationFromContext(ctx)
+	assert.NotNil(t, isolation)
+	assert.NotNil(t, isolation.persistenceServiceFacade)
+	assert.NotNil(t, isolation.gracefulShutdown)
+}
+
+func TestNewDefaultIsolation(t *testing.T) {
+	isolation := NewDefaultIsolation()
+	assert.NotNil(t, isolation)
+	assert.NotNil(t, isolation.persistenceServiceFacade)
+	assert.NotNil(t, isolation.gracefulShutdown)
+	assert.Equal(t, persistenceServiceFacade, isolation.persistenceServiceFacade)
+}

pkg/bbgo/persistence.go

@@ -10,11 +10,11 @@ import (
 	"github.com/c9s/bbgo/pkg/service"
 )
 
-var DefaultPersistenceServiceFacade = &service.PersistenceServiceFacade{
+var defaultPersistenceServiceFacade = &service.PersistenceServiceFacade{
 	Memory: service.NewMemoryService(),
 }
 
-var persistenceServiceFacade = DefaultPersistenceServiceFacade
+var persistenceServiceFacade = defaultPersistenceServiceFacade
 
 // Sync syncs the object properties into the persistence layer
 func Sync(ctx context.Context, obj interface{}) {
@@ -24,7 +24,9 @@ func Sync(ctx context.Context, obj interface{}) {
 		return
 	}
 
-	ps := persistenceServiceFacade.Get()
+	isolation := GetIsolationFromContext(ctx)
+
+	ps := isolation.persistenceServiceFacade.Get()
 	err := storePersistenceFields(obj, id, ps)
 	if err != nil {
 		log.WithError(err).Errorf("persistence sync failed")