client/isolate: correctly propagate context.Context

client/cmd/isolate/checker_test.go

 // Copyright 2016 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package main
 
 import (
         "errors"
         "fmt"
         "sync"
@@ skipping 18 matching lines @@
         itemBatches [][]*service.HandlersEndpointsV1Digest
 
         // batchc is a channel on which incoming batches are sent (before processing)
         // if non-nil.
         batchc chan<- []*service.HandlersEndpointsV1Digest
 
         // errc is a channel from which errors are pulled, if non-nil.
         errc <-chan error
 }
 
-func (f *fakeIsolateService) Contains(_ context.Context, digests []*service.HandlersEndpointsV1Digest) ([]*isolatedclient.PushState, error) {
+func (f *fakeIsolateService) Contains(ctx context.Context, digests []*service.HandlersEndpointsV1Digest) ([]*isolatedclient.PushState, error) {
+        select {
+        case <-ctx.Done():
+                return nil, ctx.Err()
+        default:
+        }
         f.mu.Lock()
         defer f.mu.Unlock()
         var states []*isolatedclient.PushState
         f.itemBatches = append(f.itemBatches, digests)
         if f.batchc != nil {
                 f.batchc <- digests
         }
         if f.pushStates == nil {
                 f.pushStates = make(map[string]*isolatedclient.PushState)
         }
@@ skipping 103 matching lines @@
                         Path:   fmt.Sprintf("/item/%d", i),
                         Digest: isolated.HexDigest(fmt.Sprintf("digest%d", i)),
                 }
                 checker.AddItem(item, false, nop)
         }
 
         if err := checker.Close(); err != errBang {
                 t.Fatalf("checker.Close: got error %v; want %v", err, errBang)
         }
 }