Initial distributor implementation

appengine/cmd/dm/mutate/schedule_execution_test.go

Index: appengine/cmd/dm/mutate/schedule_execution_test.go
diff --git a/appengine/cmd/dm/mutate/schedule_execution_test.go b/appengine/cmd/dm/mutate/schedule_execution_test.go
index c953cd7a104ac32ec33f4dac13abc285d24e8ec9..58661c962bdc5374fbe46ca9d64ad3bec2e2f26f 100644
--- a/appengine/cmd/dm/mutate/schedule_execution_test.go
+++ b/appengine/cmd/dm/mutate/schedule_execution_test.go
@@ -5,29 +5,97 @@
 package mutate
 
 import (
+	"fmt"
 	"testing"
+	"time"
 
-	"github.com/luci/gae/impl/memory"
-	"github.com/luci/luci-go/common/api/dm/service/v1"
 	. "github.com/smartystreets/goconvey/convey"
-	"golang.org/x/net/context"
+
+	"github.com/luci/gae/service/datastore"
+	"github.com/luci/luci-go/appengine/cmd/dm/distributor/fake"
+	"github.com/luci/luci-go/appengine/cmd/dm/model"
+	"github.com/luci/luci-go/common/api/dm/service/v1"
+	"github.com/luci/luci-go/common/errors"
+	. "github.com/luci/luci-go/common/testing/assertions"
 )
 
 func TestScheduleExecution(t *testing.T) {
 	t.Parallel()
 
 	Convey("ScheduleExecution", t, func() {
-		c := memory.Use(context.Background())
-		se := &ScheduleExecution{dm.NewAttemptID("quest", 1)}
+		_, c, dist, _ := fake.Setup(FinishExecutionFn)
+
+		qdesc := fake.QuestDesc("quest")
+		qid := qdesc.QuestID()
+		se := &ScheduleExecution{dm.NewAttemptID(qid, 1)}
 
 		Convey("Root", func() {
-			So(se.Root(c).String(), ShouldEqual, `dev~app::/Attempt,"quest|fffffffe"`)
+			So(se.Root(c).String(), ShouldEqual, fmt.Sprintf(`dev~app::/Attempt,"%s|fffffffe"`, qid))
 		})
 
 		Convey("RollForward", func() {
-			muts, err := se.RollForward(c)
-			So(err, ShouldBeNil)
-			So(muts, ShouldBeNil)
+			ds := datastore.Get(c)
+			q := &model.Quest{ID: qid, Desc: *qdesc}
+			a := &model.Attempt{
+				ID:    *se.For,
+				State: dm.Attempt_SCHEDULING,
+			}
+			e := model.ExecutionFromID(c, dm.NewExecutionID(qid, 1, 1))
+			So(ds.Put(q, a), ShouldBeNil)
+
+			Convey("basic", func() {
+				dist.TimeToStart = time.Minute * 5
+
+				muts, err := se.RollForward(c)
+				So(err, ShouldBeNil)
+				So(muts, ShouldBeNil)
+
+				So(ds.Get(a, e), ShouldBeNil)
+				Convey("distributor information is saved", func() {
+					tok := fake.MkToken(dm.NewExecutionID(qid, 1, 1))
+
+					So(a.State, ShouldEqual, dm.Attempt_EXECUTING)
+					So(e.State, ShouldEqual, dm.Execution_SCHEDULING)
+					So(e.DistributorConfigName, ShouldEqual, "fakeDistributor")
+					So(e.DistributorToken, ShouldEqual, tok)
+					So(e.DistributorConfigVersion, ShouldEqual, "testing")
+					So(e.TimeToStart, ShouldEqual, time.Minute*5)
+				})
+				Convey("a timeout is set", func() {
+					ex, err := ds.Exists(ds.MakeKey(
+						"Attempt", a.ID.DMEncoded(),
+						"Execution", 1,
+						"tumble.Mutation", "n:timeout"))
+					So(err, ShouldBeNil)
+					So(ex.All(), ShouldBeTrue)
+				})
+			})
+
+			Convey("transient", func() {
+				dist.RunError = errors.WrapTransient(errors.New("transient failure"))
+
+				muts, err := se.RollForward(c)
+				So(err, ShouldErrLike, "transient")
+				So(muts, ShouldBeNil)
+			})
+
+			Convey("rejection", func() {
+				dist.RunError = errors.New("no soup for you")
+
+				muts, err := se.RollForward(c)
+				So(err, ShouldBeNil)
+				So(muts, ShouldBeNil)
+
+				So(ds.Get(a, e), ShouldBeNil)
+				So(a.State, ShouldEqual, dm.Attempt_ABNORMAL_FINISHED)
+				So(a.AbnormalFinish.Status, ShouldEqual, dm.AbnormalFinish_REJECTED)
+				So(a.AbnormalFinish.Reason, ShouldContainSubstring, "non-transient")
+
+				So(e.State, ShouldEqual, dm.Execution_ABNORMAL_FINISHED)
+				So(e.AbnormalFinish.Status, ShouldEqual, dm.AbnormalFinish_REJECTED)
+				So(e.AbnormalFinish.Reason, ShouldContainSubstring, "non-transient")
+			})
+
 		})
 	})
 }