From 8ddf1348db7f44de211752c35887dc6190c24097 Mon Sep 17 00:00:00 2001
From: Brian Picciano <mediocregopher@gmail.com>
Date: Sun, 3 Feb 2019 16:05:33 -0500
Subject: [PATCH] mdb/mdatastore: refactor to use new mctx/mlog/mcfg stuff

---
 mdb/mdatastore/datastore.go      | 46 ++++++++++++++++---------
 mdb/mdatastore/datastore_test.go | 58 +++++++++++++-------------------
 2 files changed, 54 insertions(+), 50 deletions(-)

diff --git a/mdb/mdatastore/datastore.go b/mdb/mdatastore/datastore.go
index 2da96ef..04730ea 100644
--- a/mdb/mdatastore/datastore.go
+++ b/mdb/mdatastore/datastore.go
@@ -3,13 +3,12 @@
 package mdatastore
 
 import (
-	"context"
-
 	"cloud.google.com/go/datastore"
-	"github.com/mediocregopher/mediocre-go-lib/m"
-	"github.com/mediocregopher/mediocre-go-lib/mcfg"
+	"github.com/mediocregopher/mediocre-go-lib/mctx"
 	"github.com/mediocregopher/mediocre-go-lib/mdb"
+	"github.com/mediocregopher/mediocre-go-lib/merr"
 	"github.com/mediocregopher/mediocre-go-lib/mlog"
+	"github.com/mediocregopher/mediocre-go-lib/mrun"
 )
 
 // Datastore is a wrapper around a datastore client providing more
@@ -21,23 +20,38 @@ type Datastore struct {
 	log *mlog.Logger
 }
 
-// Cfg configures and returns a Datastore instance which will be usable once
-// StartRun is called on the passed in Cfg instance.
-func Cfg(cfg *mcfg.Cfg) *Datastore {
-	cfg = cfg.Child("datastore")
-	var ds Datastore
-	ds.gce = mdb.CfgGCE(cfg)
-	ds.log = m.Log(cfg, &ds)
-	cfg.Start.Then(func(ctx context.Context) error {
+// MNew returns a Datastore instance which will be initialized and configured
+// when the start event is triggered on ctx (see mrun.Start). The Datastore
+// instance will have Close called on it when the stop event is triggered on ctx
+// (see mrun.Stop).
+//
+// gce is optional and can be passed in if there's an existing gce object which
+// should be used, otherwise a new one will be created with mdb.MGCE.
+func MNew(ctx mctx.Context, gce *mdb.GCE) *Datastore {
+	if gce == nil {
+		gce = mdb.MGCE(ctx, "")
+	}
+
+	ctx = mctx.ChildOf(ctx, "datastore")
+	ds := &Datastore{
+		gce: gce,
+		log: mlog.From(ctx),
+	}
+	ds.log.SetKV(ds)
+
+	mrun.OnStart(ctx, func(innerCtx mctx.Context) error {
 		ds.log.Info("connecting to datastore")
 		var err error
-		ds.Client, err = datastore.NewClient(ctx, ds.gce.Project, ds.gce.ClientOptions()...)
-		return mlog.ErrWithKV(err, &ds)
+		ds.Client, err = datastore.NewClient(innerCtx, ds.gce.Project, ds.gce.ClientOptions()...)
+		return merr.WithKV(err, ds.KV())
 	})
-	return &ds
+	mrun.OnStop(ctx, func(mctx.Context) error {
+		return ds.Client.Close()
+	})
+	return ds
 }
 
 // KV implements the mlog.KVer interface.
-func (ds *Datastore) KV() mlog.KV {
+func (ds *Datastore) KV() map[string]interface{} {
 	return ds.gce.KV()
 }
diff --git a/mdb/mdatastore/datastore_test.go b/mdb/mdatastore/datastore_test.go
index fa7ec25..d6432c2 100644
--- a/mdb/mdatastore/datastore_test.go
+++ b/mdb/mdatastore/datastore_test.go
@@ -1,50 +1,40 @@
 package mdatastore
 
 import (
-	"context"
 	. "testing"
-	"time"
 
 	"cloud.google.com/go/datastore"
-	"github.com/mediocregopher/mediocre-go-lib/mcfg"
-	"github.com/mediocregopher/mediocre-go-lib/mdb"
 	"github.com/mediocregopher/mediocre-go-lib/mrand"
+	"github.com/mediocregopher/mediocre-go-lib/mtest"
 	"github.com/mediocregopher/mediocre-go-lib/mtest/massert"
 )
 
-var testDS *Datastore
-
-func init() {
-	mdb.DefaultGCEProject = "test"
-	cfg := mcfg.New()
-	testDS = Cfg(cfg)
-	cfg.StartTestRun()
-}
-
 // Requires datastore emulator to be running
 func TestBasic(t *T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
+	ctx := mtest.NewCtx()
+	mtest.SetEnv(ctx, "GCE_PROJECT", "test")
+	ds := MNew(ctx, nil)
+	mtest.Run(ctx, t, func() {
+		name := mrand.Hex(8)
+		key := datastore.NameKey("testKind", name, nil)
+		key.Namespace = "TestBasic_" + mrand.Hex(8)
+		type valType struct {
+			A, B int
+		}
+		val := valType{
+			A: mrand.Int(),
+			B: mrand.Int(),
+		}
 
-	name := mrand.Hex(8)
-	key := datastore.NameKey("testKind", name, nil)
-	key.Namespace = "TestBasic_" + mrand.Hex(8)
-	type valType struct {
-		A, B int
-	}
-	val := valType{
-		A: mrand.Int(),
-		B: mrand.Int(),
-	}
+		if _, err := ds.Put(ctx, key, &val); err != nil {
+			t.Fatal(err)
+		}
 
-	if _, err := testDS.Put(ctx, key, &val); err != nil {
-		t.Fatal(err)
-	}
+		var val2 valType
+		if err := ds.Get(ctx, key, &val2); err != nil {
+			t.Fatal(err)
+		}
 
-	var val2 valType
-	if err := testDS.Get(ctx, key, &val2); err != nil {
-		t.Fatal(err)
-	}
-
-	massert.Fatal(t, massert.Equal(val, val2))
+		massert.Fatal(t, massert.Equal(val, val2))
+	})
 }