Refactor distributor API so that methods always get the Quest_Desc too.

dm/appengine/distributor/fake/fake.go

 // Copyright 2015 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package fake
 
 import (
         "encoding/json"
         "fmt"
         "net/http"
@@ skipping 49 matching lines @@
                         "acls.cfg": `
                         readers: "reader_group"
                         writers: "writer_group"
                         `,
                 },
         }))
         c = auth.WithState(c, &authtest.FakeState{
                 Identity: identity.AnonymousIdentity,
         })
         dist = &Distributor{}
-        reg := distributor.NewTestingRegistry(map[string]distributor.D{
-                "fakeDistributor": dist,
+        reg := distributor.NewTestingRegistry(distributor.TestFactoryMap{
+                "fakeDistributor": func(c context.Context, cfg *distributor.Config) distributor.D {
+                        return &BoundDistributor{dist, c, cfg}
+                },
         }, fn)
         c = distributor.WithRegistry(c, reg)
         return
 }
 
 // DistributorData is the blob of data that the fake.Distributor keeps when DM
 // calls its Run method. This is roughly equivalent to the state that
 // a distributor (like swarming) would store in its own datastore about a job.
 type DistributorData struct {
         NotifyTopic pubsub.Topic
@@ skipping 139 matching lines @@
 type Distributor struct {
         // RunError can be set to make Run return this error when it's invoked.
         RunError error
         // This can be set to turn the distributor into a polling-based distributor.
         PollbackTime time.Duration
 
         sync.Mutex
         tasks map[distributor.Token]*DistributorData
 }
 
+// BoundDistributor binds the fake.Distributor to a Context and
+// a distributor.Config. It implements distributor.D.
+type BoundDistributor struct {
+        *Distributor
+        c   context.Context
+        cfg *distributor.Config
+}
+
 // MkToken makes a distributor Token out of an Execution_ID. In this
 // implementation of a Distributor there's a 1:1 mapping between Execution_ID
 // and distributor task. This is not always the case for real distributor
 // implementations.
 func MkToken(eid *dm.Execution_ID) distributor.Token {
         return distributor.Token(fmt.Sprintf("fakeDistributor:%s|%d|%d", eid.Quest,
                 eid.Attempt, eid.Id))
 }
 
 // Run implements distributor.D
-func (f *Distributor) Run(desc *distributor.TaskDescription) (tok distributor.Token, pollbackTime time.Duration, err error) {
-        if err = f.RunError; err != nil {
+func (d *BoundDistributor) Run(desc *dm.Quest_Desc, exAuth *dm.Execution_Auth, prev *dm.JsonResult) (tok distributor.Token, pollbackTime time.Duration, err error) {
+        if err = d.RunError; err != nil {
                 return
         }
-        pollbackTime = f.PollbackTime
+        pollbackTime = d.PollbackTime
 
-        exAuth := desc.ExecutionAuth()
         tok = MkToken(exAuth.Id)

[Vadim Sh., 2016/09/20 00:24:26]

is it really that difficult to type two more letters to make it "MakeToken"? :)
I'm just curious why...

[iannucci, 2016/09/20 00:51:28]

I think it's fine for a testing library... Mk is a pretty standard abbreviation.
I guess I could rename it CreateDistributorTokenFromExecutionID. When is an
abbreviation not OK?

 
         tsk := &DistributorData{
                 Auth:  exAuth,
-                Desc:  desc.Payload(),
-                State: desc.PreviousResult(),
+                Desc:  desc,
+                State: prev,
         }
-        tsk.NotifyTopic, tsk.NotifyAuth, err = desc.PrepareTopic()
+        tsk.NotifyTopic, tsk.NotifyAuth, err = d.cfg.PrepareTopic(d.c, exAuth.Id)
         panicIf(err)
 
-        f.Lock()
-        defer f.Unlock()
-        if f.tasks == nil {
-                f.tasks = map[distributor.Token]*DistributorData{}
+        d.Lock()
+        defer d.Unlock()
+        if d.tasks == nil {
+                d.tasks = map[distributor.Token]*DistributorData{}
         }
-        f.tasks[tok] = tsk
+        d.tasks[tok] = tsk
         return
 }
 
 // Cancel implements distributor.D
-func (f *Distributor) Cancel(tok distributor.Token) (err error) {
-        f.Lock()
-        defer f.Unlock()
-        if tsk, ok := f.tasks[tok]; ok {
+func (d *BoundDistributor) Cancel(_ *dm.Quest_Desc, tok distributor.Token) (err error) {
+        d.Lock()
+        defer d.Unlock()
+        if tsk, ok := d.tasks[tok]; ok {
                 tsk.done = true
                 tsk.abnorm = &dm.AbnormalFinish{
                         Status: dm.AbnormalFinish_CANCELLED,
                         Reason: "cancelled via Cancel()"}
         } else {
                 err = fmt.Errorf("MISSING task %q", tok)
         }
         return
 }
 
 // GetStatus implements distributor.D
-func (f *Distributor) GetStatus(tok distributor.Token) (rslt *dm.Result, err error) {
-        f.Lock()
-        defer f.Unlock()
-        if tsk, ok := f.tasks[tok]; ok {
+func (d *BoundDistributor) GetStatus(_ *dm.Quest_Desc, tok distributor.Token) (rslt *dm.Result, err error) {
+        d.Lock()
+        defer d.Unlock()
+        if tsk, ok := d.tasks[tok]; ok {
                 if tsk.done {
                         if tsk.abnorm != nil {
                                 rslt = &dm.Result{AbnormalFinish: tsk.abnorm}
                         } else {
                                 rslt = &dm.Result{Data: tsk.State}
                         }
                 }
         } else {
                 rslt = &dm.Result{
                         AbnormalFinish: &dm.AbnormalFinish{
                                 Status: dm.AbnormalFinish_MISSING,
                                 Reason: fmt.Sprintf("unknown token: %s", tok)},
                 }
         }
         return
 }
 
 // FakeURLPrefix is the url that all fake InfoURLs are prefixed with.
 const FakeURLPrefix = "https://info.example.com/"
 
 // InfoURL builds a fake InfoURL for the given Execution_ID
 func InfoURL(e *dm.Execution_ID) string {
         return FakeURLPrefix + string(MkToken(e))
 }
 
 // InfoURL implements distributor.D
-func (f *Distributor) InfoURL(tok distributor.Token) string {
+func (d *BoundDistributor) InfoURL(tok distributor.Token) string {
         return FakeURLPrefix + string(tok)
 }
 
 // HandleNotification implements distributor.D
-func (f *Distributor) HandleNotification(n *distributor.Notification) (rslt *dm.Result, err error) {
-        return f.GetStatus(distributor.Token(n.Attrs["token"]))
+func (d *BoundDistributor) HandleNotification(q *dm.Quest_Desc, n *distributor.Notification) (rslt *dm.Result, err error) {
+        return d.GetStatus(q, distributor.Token(n.Attrs["token"]))
 }
 
 // HandleTaskQueueTask is not implemented, and shouldn't be needed for most
 // tests. It could be implemented if some new test required it, however.
-func (f *Distributor) HandleTaskQueueTask(r *http.Request) ([]*distributor.Notification, error) {
+func (d *BoundDistributor) HandleTaskQueueTask(r *http.Request) ([]*distributor.Notification, error) {
         panic("not implemented")
 }
 
 // Validate implements distributor.D (by returning a nil error for every
 // payload).
-func (f *Distributor) Validate(payload string) error {
+func (d *BoundDistributor) Validate(payload string) error {
         return nil
 }
 
 // RunTask allows you to run the task associated with the provided execution id.
 //
 // If the task corresponding to `eid` returns an error, or if the distributor
 // itself actually has an error, this method will return an error. Notably, if
 // `cb` returns an error, it will simply mark the corresponding task as FAILED,
 // but will return nil here.
 //
 // If the task exists and hasn't been run yet, cb will be called, and can do
 // anything that you may want to a test to do. Think of the callback as the
 // recipe engine; it has the opportunity to do anything it wants to, interact
 // with the deps server (or not), succeed (or not), etc.
 //
 // If the callback needs to maintain state between executions, Task.State is
 // read+write; when the callback exits, the final value of Task.State will be
 // passed back to the DM instance under test. A re-execution of the attempt will
 // start with the new value.
-func (f *Distributor) RunTask(c context.Context, eid *dm.Execution_ID, cb func(*Task) error) (err error) {
+func (d *Distributor) RunTask(c context.Context, eid *dm.Execution_ID, cb func(*Task) error) (err error) {
         tok := MkToken(eid)
 
-        f.Lock()
-        tsk := f.tasks[tok]
+        d.Lock()
+        tsk := d.tasks[tok]
         if tsk == nil {
                 err = fmt.Errorf("cannot RunTask(%q): doesn't exist", tok)
         } else {
                 if tsk.done {
                         err = fmt.Errorf("cannot RunTask(%q): running twice", tok)
                 } else {
                         tsk.done = true
                 }
         }
-        f.Unlock()
+        d.Unlock()
 
         if err != nil {
                 return
         }
 
         abnorm := (*dm.AbnormalFinish)(nil)
 
         usrTsk := &Task{
                 tsk.Auth,
                 tsk.Desc,
                 tsk.State,
         }
 
         defer func() {
-                f.Lock()
+                d.Lock()
                 {
                         tsk.abnorm = abnorm
                         tsk.State = usrTsk.State
 
                         if r := recover(); r != nil {
                                 tsk.abnorm = &dm.AbnormalFinish{
                                         Status: dm.AbnormalFinish_CRASHED,
                                         Reason: fmt.Sprintf("caught panic: %q", r),
                                 }
                         }
                 }
-                f.Unlock()
+                d.Unlock()
 
                 err = tumble.RunMutation(c, &distributor.NotifyExecution{
                         CfgName: "fakeDistributor",
                         Notification: &distributor.Notification{
                                 ID:    tsk.Auth.Id,
                                 Attrs: map[string]string{"token": string(tok)}},
                 })
         }()
 
         err = cb(usrTsk)
         if err != nil {
                 err = nil
                 abnorm = &dm.AbnormalFinish{
                         Status: dm.AbnormalFinish_FAILED,
                         Reason: fmt.Sprintf("cb error: %q", err),
                 }
         }
         return
 }
 
 func panicIf(err error) {
         if err != nil {
                 panic(err)
         }
 }
 
-var _ distributor.D = (*Distributor)(nil)
+var _ distributor.D = (*BoundDistributor)(nil)
 
 // QuestDesc generates a normalized generic QuestDesc of the form:
 //   Quest_Desc{
 //     DistributorConfigName: "fakeDistributor",
 //     Parameters:            `{"name":"$name"}`,
 //     DistributorParameters: "{}",
 //   }
 func QuestDesc(name string) *dm.Quest_Desc {
         params, err := json.Marshal(struct {
                 Name string `json:"name"`
@@ skipping 46 matching lines @@
                         panic(fmt.Errorf("expected 1 arg on rhs, got %d", len(expect)))
                 }
                 e := expect[0].(*dm.GraphData)
                 ret, err := s.WalkGraph(c, r)
                 if nilExpect := assertions.ShouldErrLike(err, nil); nilExpect != "" {
                         return nilExpect
                 }
                 return convey.ShouldResemble(normalize(ret), e)
         }
 }