Support unit testing Perform funcs which push jobs

context.go

@@ -13,8 +13,7 @@ import (
 type valueKey int
 
 const (
-	poolKey valueKey = 2
-	jobKey  valueKey = 3
+	helperKey valueKey = 1
 )
 
 var (
@@ -57,6 +56,10 @@ type Helper interface {
 	// allows direct access to the Faktory server from the job
 	With(func(*faktory.Client) error) error
 
+	// Push the given job to Faktory. In test mode, this will
+	// save the job in an array so you can assert it.
+	Push(*faktory.Job) error
+
 	// Faktory Enterprise:
 	// this method integrates with Faktory Enterprise's Job Tracking feature.
 	// `reserveUntil` is optional, only needed for long jobs which have more dynamic
@@ -68,6 +71,47 @@ type Helper interface {
 	TrackProgress(percent int, desc string, reserveUntil *time.Time) error
 }
 
+// The test helper is a mode of operation which tries to avoid network calls.
+// This allows you to unit test your job code with the Helper API without
+// Faktory running.
+type testHelper struct {
+	job        *faktory.Job
+	pushedJobs []*faktory.Job
+}
+
+func (h *testHelper) Jid() string {
+	return h.job.Jid
+}
+func (h *testHelper) Bid() string {
+	if b, ok := h.job.GetCustom("_bid"); ok {
+		return b.(string)
+	}
+	return ""
+}
+func (h *testHelper) JobType() string {
+	return h.job.Type
+}
+
+func (h *testHelper) TrackProgress(percent int, desc string, reserveUntil *time.Time) error {
+	return nil
+}
+func (h *testHelper) Batch(fn func(*faktory.Batch) error) error {
+	return nil
+}
+func (h *testHelper) With(fn func(*faktory.Client) error) error {
+	return fmt.Errorf("With() does not work in test mode, you will need to abstract out this call")
+}
+
+func (h *testHelper) Push(job *faktory.Job) error {
+	h.pushedJobs = append(h.pushedJobs, job)
+	return nil
+}
+
+func (h *testHelper) PushedJobs() []*faktory.Job {
+	return h.pushedJobs
+}
+
+// This "live" helper impl requires Faktory to be running.
 type jobHelper struct {
 	job  *faktory.Job
 	pool *faktory.Pool
@@ -90,24 +134,29 @@ func (h *jobHelper) JobType() string {
 // context of an executing job. It will panic if it cannot
 // create a Helper due to missing context values.
 func HelperFor(ctx context.Context) Helper {
-	if j := ctx.Value(jobKey); j != nil {
-		job := j.(*faktory.Job)
-		if p := ctx.Value(poolKey); p != nil {
-			pool := p.(*faktory.Pool)
-			return &jobHelper{
-				job:  job,
-				pool: pool,
-			}
-		}
+	if j := ctx.Value(helperKey); j != nil {
+		helper := j.(Helper)
+		return helper
 	}
 	log.Panic("Invalid job context, cannot create faktory_worker_go job helper")
 	return nil
 }
 
-func jobContext(pool *faktory.Pool, job *faktory.Job) context.Context {
+func jobContext(testing bool, pool *faktory.Pool, job *faktory.Job) context.Context {
+	var helper Helper
+	if testing {
+		helper = &testHelper{
+			job: job,
+		}
+	} else {
+		helper = &jobHelper{
+			job:  job,
+			pool: pool,
+		}
+	}
+
 	ctx := context.Background()
-	ctx = context.WithValue(ctx, poolKey, pool)
-	ctx = context.WithValue(ctx, jobKey, job)
+	ctx = context.WithValue(ctx, helperKey, helper)
 	return ctx
 }
 
@@ -146,3 +195,9 @@ func (h *jobHelper) Batch(fn func(*faktory.Batch) error) error {
 func (h *jobHelper) With(fn func(*faktory.Client) error) error {
 	return h.pool.With(fn)
 }
+
+func (h *jobHelper) Push(job *faktory.Job) error {
+	return h.pool.With(func(cl *faktory.Client) error {
+		return cl.Push(job)
+	})
+}

context_test.go

@@ -22,7 +22,7 @@ func TestSimpleContext(t *testing.T) {
 	//assert.NoError(t, err)
 	//cl.Push(job)
 
-	ctx := jobContext(pool, job)
+	ctx := jobContext(false, pool, job)
 	help := HelperFor(ctx)
 	assert.Equal(t, help.Jid(), job.Jid)
 	assert.Empty(t, help.Bid())
@@ -51,7 +51,7 @@ func TestBatchContext(t *testing.T) {
 	job.SetCustom("track", 1)
 	job.SetCustom("_bid", "nosuchbatch")
 
-	ctx := jobContext(pool, job)
+	ctx := jobContext(false, pool, job)
 	help := HelperFor(ctx)
 	assert.Equal(t, help.Jid(), job.Jid)
 	assert.Equal(t, "nosuchbatch", help.Bid())

middleware_test.go

@@ -31,7 +31,7 @@ func TestMiddleware(t *testing.T) {
 	}
 
 	job := faktory.NewJob("blah", 1, 2)
-	ctx := jobContext(mgr.Pool, job)
+	ctx := jobContext(false, mgr.Pool, job)
 	assert.Nil(t, ctx.Value("a"))
 	assert.EqualValues(t, 0, counter)
 

runner.go

@@ -125,7 +125,7 @@ func processOne(mgr *Manager) error {
 		return je
 	}
 
-	joberr := dispatch(mgr.middleware, jobContext(mgr.Pool, job), job, perform)
+	joberr := dispatch(mgr.middleware, jobContext(false, mgr.Pool, job), job, perform)
 	if joberr != nil {
 		// job errors are normal and expected, we don't return early from them
 		mgr.Logger.Errorf("Error running %s job %s: %v", job.Type, job.Jid, joberr)

testing.go

@@ -1,13 +1,20 @@
 package faktory_worker
 
 import (
+	"context"
 	"encoding/json"
 
 	faktory "github.com/contribsys/faktory/client"
 )
 
 type PerformExecutor interface {
+	// This executes the Perform func with full network
+	// access to the helper.
 	Execute(*faktory.Job, Perform) error
+
+	// This stubs out the helper so your Perform func can
+	// use it without network access.
+	ExecuteWithResult(*faktory.Job, Perform) (ExecuteResult, error)
 }
 
 type testExecutor struct {
@@ -19,18 +26,35 @@ func NewTestExecutor(p *faktory.Pool) PerformExecutor {
 }
 
 func (tp *testExecutor) Execute(specjob *faktory.Job, p Perform) error {
+	_, err := tp.exec(specjob, p, false)
+	return err
+}
+
+func (tp *testExecutor) ExecuteWithResult(specjob *faktory.Job, p Perform) (ExecuteResult, error) {
+	ctx, err := tp.exec(specjob, p, true)
+	if err != nil {
+		return nil, err
+	}
+	return HelperFor(ctx).(ExecuteResult), nil
+}
+
+type ExecuteResult interface {
+	PushedJobs() []*faktory.Job
+}
+
+func (tp *testExecutor) exec(specjob *faktory.Job, p Perform, local bool) (context.Context, error) {
 	// perform a JSON round trip to ensure Perform gets the arguments
 	// exactly how a round trip to Faktory would look.
 	data, err := json.Marshal(specjob)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	var job faktory.Job
 	err = json.Unmarshal(data, &job)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	c := jobContext(tp.Pool, &job)
-	return p(c, job.Args...)
+	c := jobContext(local, tp.Pool, &job)
+	return c, p(c, job.Args...)
 }

testing_test.go

@@ -31,9 +31,13 @@ func TestTesting(t *testing.T) {
 
 	ajob := faktory.NewJob("sometype", 12, "foobar")
 	err = perf.Execute(ajob, func(ctx context.Context, args ...interface{}) error {
+		helper := HelperFor(ctx)
 		assert.Equal(t, 2, len(args))
 		assert.EqualValues(t, 12, args[0])
 		assert.EqualValues(t, "foobar", args[1])
+		assert.EqualValues(t, ajob.Jid, helper.Jid())
+		assert.EqualValues(t, "", helper.Bid())
+		assert.EqualValues(t, ajob.Type, helper.JobType())
 		return nil
 	})
 	assert.NoError(t, err)
@@ -42,3 +46,31 @@ func TestTesting(t *testing.T) {
 	})
 	assert.Equal(t, "Oops", err.Error())
 }
+
+func TestRecursiveTest(t *testing.T) {
+	pool, err := faktory.NewPool(5)
+	assert.NoError(t, err)
+	perf := NewTestExecutor(pool)
+
+	somejob := faktory.NewJob("pushtype", 12, "foobar")
+	res, err := perf.ExecuteWithResult(somejob, func(ctx context.Context, args ...interface{}) error {
+		helper := HelperFor(ctx)
+		helper.TrackProgress(10, "Starting", nil)
+
+		assert.NotEmpty(t, helper.Jid())
+		assert.Empty(t, helper.Bid())
+		assert.Equal(t, "pushtype", helper.JobType())
+
+		ajob := faktory.NewJob("sometype", 12, "foobar")
+		ajob.Queue = "whale"
+		helper.Push(ajob)
+
+		helper.TrackProgress(100, "Finished", nil)
+		return nil
+	})
+
+	assert.NoError(t, err)
+	assert.NotNil(t, res)
+	assert.EqualValues(t, 1, len(res.PushedJobs()))
+	assert.EqualValues(t, "whale", res.PushedJobs()[0].Queue)
+}