func setup(s *testerator.Setup) error {
if sniff, ok := s.Context.Value(&ctxKey).(*searchSniffer); ok {
sniff.searchIndexDocumentRequests = nil
return nil
}
sniff := &searchSniffer{}
s.Context = context.WithValue(s.Context, &ctxKey, sniff)
s.Context = appengine.WithAPICallFunc(s.Context, func(ctx context.Context, service, method string, in, out proto.Message) error {
if service == "search" && method == "IndexDocument" {
b, err := proto.Marshal(in)
if err != nil {
return err
}
req := &searchpb.IndexDocumentRequest{}
err = proto.Unmarshal(b, req)
if err != nil {
return err
}
sniff.searchIndexDocumentRequests = append(sniff.searchIndexDocumentRequests, req)
}
return appengine.APICall(ctx, service, method, in, out)
})
return nil
}
func setup(ctx context.Context, method, path, query string, header http.Header) context.Context {
stats := &requestStats{
Method: method,
Path: path,
Query: query,
Start: time.Now(),
}
if u := user.Current(ctx); u != nil {
stats.User = u.String()
stats.Admin = u.Admin
}
ctx = context.WithValue(ctx, statsKey, stats)
ctx = context.WithValue(ctx, headerKey, header)
ctx = appengine.WithAPICallFunc(ctx, appengine.APICallFunc(override))
return ctx
}
// WithContext enables profiling of functions without a corresponding request,
// as in the appengine/delay package. method and path may be empty.
func WithContext(ctx context.Context, method, path string, f func(context.Context)) {
stats := &requestStats{
Method: method,
Path: path,
Start: time.Now(),
}
if u := user.Current(ctx); u != nil {
stats.User = u.String()
stats.Admin = u.Admin
}
ctx = context.WithValue(ctx, statsKey, stats)
ctx = appengine.WithAPICallFunc(ctx, override)
f(ctx)
save(ctx)
}
// newContext creates a new timing-aware context from req.
func newContext(r *http.Request) context.Context {
ctx := appengine.NewContext(r)
stats := &requestStats{
Method: r.Method,
Path: r.URL.Path,
Query: r.URL.RawQuery,
Start: time.Now(),
}
if u := user.Current(ctx); u != nil {
stats.User = u.String()
stats.Admin = u.Admin
}
ctx = context.WithValue(ctx, statsKey, stats)
ctx = context.WithValue(ctx, headerKey, r.Header)
ctx = appengine.WithAPICallFunc(ctx, override)
return ctx
}