X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/25be2534e8746475ddc799042fefa08bd0548e9d..f81427c5f93c159d26c3125aaa7ede4c5986dd07:/services/keepstore/azure_blob_volume_test.go

diff --git a/services/keepstore/azure_blob_volume_test.go b/services/keepstore/azure_blob_volume_test.go
index 5d556b3e8c..232382c421 100644
--- a/services/keepstore/azure_blob_volume_test.go
+++ b/services/keepstore/azure_blob_volume_test.go
@@ -2,13 +2,13 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"crypto/md5"
 	"encoding/base64"
 	"encoding/xml"
 	"flag"
 	"fmt"
 	"io/ioutil"
-	"log"
 	"math/rand"
 	"net"
 	"net/http"
@@ -21,6 +21,7 @@ import (
 	"testing"
 	"time"
 
+	log "github.com/Sirupsen/logrus"
 	"github.com/curoverse/azure-sdk-for-go/storage"
 )
 
@@ -365,7 +366,13 @@ func NewTestableAzureBlobVolume(t TB, readonly bool, replication int) *TestableA
 		}
 	}
 
-	v := NewAzureBlobVolume(azClient, container, readonly, replication)
+	v := &AzureBlobVolume{
+		ContainerName:    container,
+		ReadOnly:         readonly,
+		AzureReplication: replication,
+		azClient:         azClient,
+		bsClient:         azClient.GetBlobService(),
+	}
 
 	return &TestableAzureBlobVolume{
 		AzureBlobVolume: v,
@@ -448,12 +455,12 @@ func TestAzureBlobVolumeRangeFenceposts(t *testing.T) {
 			data[i] = byte((i + 7) & 0xff)
 		}
 		hash := fmt.Sprintf("%x", md5.Sum(data))
-		err := v.Put(hash, data)
+		err := v.Put(context.Background(), hash, data)
 		if err != nil {
 			t.Error(err)
 		}
 		gotData := make([]byte, len(data))
-		gotLen, err := v.Get(hash, gotData)
+		gotLen, err := v.Get(context.Background(), hash, gotData)
 		if err != nil {
 			t.Error(err)
 		}
@@ -494,7 +501,7 @@ func TestAzureBlobVolumeCreateBlobRace(t *testing.T) {
 	allDone := make(chan struct{})
 	v.azHandler.race = make(chan chan struct{})
 	go func() {
-		err := v.Put(TestHash, TestBlock)
+		err := v.Put(context.Background(), TestHash, TestBlock)
 		if err != nil {
 			t.Error(err)
 		}
@@ -504,7 +511,7 @@ func TestAzureBlobVolumeCreateBlobRace(t *testing.T) {
 	v.azHandler.race <- continuePut
 	go func() {
 		buf := make([]byte, len(TestBlock))
-		_, err := v.Get(TestHash, buf)
+		_, err := v.Get(context.Background(), TestHash, buf)
 		if err != nil {
 			t.Error(err)
 		}
@@ -547,7 +554,7 @@ func TestAzureBlobVolumeCreateBlobRaceDeadline(t *testing.T) {
 	go func() {
 		defer close(allDone)
 		buf := make([]byte, BlockSize)
-		n, err := v.Get(TestHash, buf)
+		n, err := v.Get(context.Background(), TestHash, buf)
 		if err != nil {
 			t.Error(err)
 			return
@@ -569,12 +576,76 @@ func TestAzureBlobVolumeCreateBlobRaceDeadline(t *testing.T) {
 	}
 }
 
+func TestAzureBlobVolumeContextCancelGet(t *testing.T) {
+	testAzureBlobVolumeContextCancel(t, func(ctx context.Context, v *TestableAzureBlobVolume) error {
+		v.PutRaw(TestHash, TestBlock)
+		_, err := v.Get(ctx, TestHash, make([]byte, BlockSize))
+		return err
+	})
+}
+
+func TestAzureBlobVolumeContextCancelPut(t *testing.T) {
+	testAzureBlobVolumeContextCancel(t, func(ctx context.Context, v *TestableAzureBlobVolume) error {
+		return v.Put(ctx, TestHash, make([]byte, BlockSize))
+	})
+}
+
+func TestAzureBlobVolumeContextCancelCompare(t *testing.T) {
+	testAzureBlobVolumeContextCancel(t, func(ctx context.Context, v *TestableAzureBlobVolume) error {
+		v.PutRaw(TestHash, TestBlock)
+		return v.Compare(ctx, TestHash, TestBlock2)
+	})
+}
+
+func testAzureBlobVolumeContextCancel(t *testing.T, testFunc func(context.Context, *TestableAzureBlobVolume) error) {
+	defer func(t http.RoundTripper) {
+		http.DefaultTransport = t
+	}(http.DefaultTransport)
+	http.DefaultTransport = &http.Transport{
+		Dial: (&azStubDialer{}).Dial,
+	}
+
+	v := NewTestableAzureBlobVolume(t, false, 3)
+	defer v.Teardown()
+	v.azHandler.race = make(chan chan struct{})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	allDone := make(chan struct{})
+	go func() {
+		defer close(allDone)
+		err := testFunc(ctx, v)
+		if err != context.Canceled {
+			t.Errorf("got %T %q, expected %q", err, err, context.Canceled)
+		}
+	}()
+	releaseHandler := make(chan struct{})
+	select {
+	case <-allDone:
+		t.Error("testFunc finished without waiting for v.azHandler.race")
+	case <-time.After(10 * time.Second):
+		t.Error("timed out waiting to enter handler")
+	case v.azHandler.race <- releaseHandler:
+	}
+
+	cancel()
+
+	select {
+	case <-time.After(10 * time.Second):
+		t.Error("timed out waiting to cancel")
+	case <-allDone:
+	}
+
+	go func() {
+		<-releaseHandler
+	}()
+}
+
 func (v *TestableAzureBlobVolume) PutRaw(locator string, data []byte) {
-	v.azHandler.PutRaw(v.containerName, locator, data)
+	v.azHandler.PutRaw(v.ContainerName, locator, data)
 }
 
 func (v *TestableAzureBlobVolume) TouchWithDate(locator string, lastPut time.Time) {
-	v.azHandler.TouchWithDate(v.containerName, locator, lastPut)
+	v.azHandler.TouchWithDate(v.ContainerName, locator, lastPut)
 }
 
 func (v *TestableAzureBlobVolume) Teardown() {