1 files changed, 53 insertions, 0 deletions
diff --git a/vendor/google.golang.org/appengine/delay/delay_test.go b/vendor/google.golang.org/appengine/delay/delay_test.go
index 1c37e79..3df2bf7 100644
--- a/vendor/google.golang.org/appengine/delay/delay_test.go
+++ b/vendor/google.golang.org/appengine/delay/delay_test.go
@@ -94,6 +94,14 @@ var (
 			dupeWhich = 2
 		}
 	})
+
+	reqFuncRuns    = 0
+	reqFuncHeaders *taskqueue.RequestHeaders
+	reqFuncErr     error
+	reqFunc        = Func("req", func(c context.Context) {
+		reqFuncRuns++
+		reqFuncHeaders, reqFuncErr = RequestHeaders(c)
+	})
 )
 
 type fakeContext struct {
@@ -373,3 +381,48 @@ func TestDuplicateFunction(t *testing.T) {
 		t.Errorf("dupeWhich = %d; want 2", dupeWhich)
 	}
 }
+
+func TestGetRequestHeadersFromContext(t *testing.T) {
+	c := newFakeContext()
+
+	// Outside a delay.Func should return an error.
+	headers, err := RequestHeaders(c.ctx)
+	if headers != nil {
+		t.Errorf("RequestHeaders outside Func, got %v, want nil", headers)
+	}
+	if err != errOutsideDelayFunc {
+		t.Errorf("RequestHeaders outside Func err, got %v, want %v", err, errOutsideDelayFunc)
+	}
+
+	// Fake out the adding of a task.
+	var task *taskqueue.Task
+	taskqueueAdder = func(_ context.Context, tk *taskqueue.Task, queue string) (*taskqueue.Task, error) {
+		if queue != "" {
+			t.Errorf(`Got queue %q, expected ""`, queue)
+		}
+		task = tk
+		return tk, nil
+	}
+
+	reqFunc.Call(c.ctx)
+
+	reqFuncRuns, reqFuncHeaders = 0, nil // reset state
+	// Simulate the Task Queue service.
+	req, err := http.NewRequest("POST", path, bytes.NewBuffer(task.Payload))
+	req.Header.Set("x-appengine-taskname", "foobar")
+	if err != nil {
+		t.Fatalf("Failed making http.Request: %v", err)
+	}
+	rw := httptest.NewRecorder()
+	runFunc(c.ctx, rw, req)
+
+	if reqFuncRuns != 1 {
+		t.Errorf("reqFuncRuns: got %d, want 1", reqFuncRuns)
+	}
+	if reqFuncHeaders.TaskName != "foobar" {
+		t.Errorf("reqFuncHeaders.TaskName: got %v, want 'foobar'", reqFuncHeaders.TaskName)
+	}
+	if reqFuncErr != nil {
+		t.Errorf("reqFuncErr: got %v, want nil", reqFuncErr)
+	}
+}