Fix potential race in capabilities test.

capabilities.go

@@ -48,9 +48,6 @@ const (
 	maxInvalidateInterval = time.Minute
 )
 
-// Can be overwritten by tests.
-var getCapabilitiesNow = time.Now
-
 type capabilitiesEntry struct {
 	nextUpdate   time.Time
 	capabilities map[string]interface{}
@@ -59,6 +56,9 @@ type capabilitiesEntry struct {
 type Capabilities struct {
 	mu sync.RWMutex
 
+	// Can be overwritten by tests.
+	getNow func() time.Time
+
 	version        string
 	pool           *HttpClientPool
 	entries        map[string]*capabilitiesEntry
@@ -67,6 +67,8 @@ type Capabilities struct {
 
 func NewCapabilities(version string, pool *HttpClientPool) (*Capabilities, error) {
 	result := &Capabilities{
+		getNow: time.Now,
+
 		version:        version,
 		pool:           pool,
 		entries:        make(map[string]*capabilitiesEntry),
@@ -94,7 +96,7 @@ func (c *Capabilities) getCapabilities(key string) (map[string]interface{}, bool
 	c.mu.RLock()
 	defer c.mu.RUnlock()
 
-	now := getCapabilitiesNow()
+	now := c.getNow()
 	if entry, found := c.entries[key]; found && entry.nextUpdate.After(now) {
 		return entry.capabilities, true
 	}
@@ -103,22 +105,23 @@ func (c *Capabilities) getCapabilities(key string) (map[string]interface{}, bool
 }
 
 func (c *Capabilities) setCapabilities(key string, capabilities map[string]interface{}) {
-	now := getCapabilitiesNow()
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	now := c.getNow()
 	entry := &capabilitiesEntry{
 		nextUpdate:   now.Add(CapabilitiesCacheDuration),
 		capabilities: capabilities,
 	}
 
-	c.mu.Lock()
-	defer c.mu.Unlock()
 	c.entries[key] = entry
 }
 
 func (c *Capabilities) invalidateCapabilities(key string) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	now := getCapabilitiesNow()
+	now := c.getNow()
 	if entry, found := c.nextInvalidate[key]; found && entry.After(now) {
 		return
 	}

capabilities_test.go

@@ -120,13 +120,20 @@ func NewCapabilitiesForTest(t *testing.T) (*url.URL, *Capabilities) {
 	return NewCapabilitiesForTestWithCallback(t, nil)
 }
 
-func SetCapabilitiesGetNow(t *testing.T, f func() time.Time) {
-	old := getCapabilitiesNow
+func SetCapabilitiesGetNow(t *testing.T, capabilities *Capabilities, f func() time.Time) {
+	capabilities.mu.Lock()
+	defer capabilities.mu.Unlock()
+
+	old := capabilities.getNow
+
 	t.Cleanup(func() {
-		getCapabilitiesNow = old
+		capabilities.mu.Lock()
+		defer capabilities.mu.Unlock()
+
+		capabilities.getNow = old
 	})
 
-	getCapabilitiesNow = f
+	capabilities.getNow = f
 }
 
 func TestCapabilities(t *testing.T) {
@@ -248,7 +255,7 @@ func TestInvalidateCapabilities(t *testing.T) {
 	}
 
 	// At a later time, invalidating can be done again.
-	SetCapabilitiesGetNow(t, func() time.Time {
+	SetCapabilitiesGetNow(t, capabilities, func() time.Time {
 		return time.Now().Add(2 * time.Minute)
 	})