Better attempt at an appengine wrapper.

go/src/infra/gae/libs/wrapper/mathrand.go

+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package wrapper
+
+import (
+        "math/rand"
+
+        "golang.org/x/net/context"
+)
+
+// MathRandFactory is the function signature for factory methods compatible with
+// SetMathRandFactory.
+type MathRandFactory func(context.Context) *rand.Rand
+
+// GetMathRand gets a *"math/rand".Rand from the context. If one hasn't been
+// set, this creates a new Rand object with a Source initialized from the
+// current time accordint to GetTimeNow(c).UnixNano().
+func GetMathRand(c context.Context) *rand.Rand {
+        obj := c.Value(mathRandKey)
+        if obj == nil || obj.(MathRandFactory) == nil {
+                return rand.New(rand.NewSource(GetTimeNow(c).UnixNano()))
+        }
+        return obj.(MathRandFactory)(c)
+}
+
+// SetMathRandFactory sets the function to produce *"math/rand".Rand instances,
+// as returned by the GetMathRand method.
+func SetMathRandFactory(c context.Context, mrf MathRandFactory) context.Context {
+        return context.WithValue(c, mathRandKey, mrf)
+}
+
+// SetMathRand sets the current *"math/rand".Rand object in the context. Useful
+// for testing with a quick mock. This is just a shorthand SetMathRandFactory
+// invocation to set a factory which always returns the same object.
+func SetMathRand(c context.Context, r *rand.Rand) context.Context {
+        f := func(context.Context) *rand.Rand { return r }
+        if r == nil {
+                f = nil
+        }
+        return SetMathRandFactory(c, f)
+}