From 695bac7b0f4891b613b2e4e2cfd2bdd798f89b7d Mon Sep 17 00:00:00 2001
From: Brian Bland <brian.bland@docker.com>
Date: Mon, 1 Feb 2016 15:34:36 -0800
Subject: [PATCH] Adds test for S3 storage class configuration option

Signed-off-by: Brian Bland <brian.bland@docker.com>
---
 registry/storage/driver/s3/s3_test.go | 74 ++++++++++++++++++++++++---
 1 file changed, 67 insertions(+), 7 deletions(-)

diff --git a/registry/storage/driver/s3/s3_test.go b/registry/storage/driver/s3/s3_test.go
index be099790..660d5350 100644
--- a/registry/storage/driver/s3/s3_test.go
+++ b/registry/storage/driver/s3/s3_test.go
@@ -18,7 +18,7 @@ import (
 // Hook up gocheck into the "go test" runner.
 func Test(t *testing.T) { check.TestingT(t) }
 
-var s3DriverConstructor func(rootDirectory string) (*Driver, error)
+var s3DriverConstructor func(rootDirectory string, storageClass s3.StorageClass) (*Driver, error)
 var skipS3 func() string
 
 func init() {
@@ -35,7 +35,7 @@ func init() {
 	}
 	defer os.Remove(root)
 
-	s3DriverConstructor = func(rootDirectory string) (*Driver, error) {
+	s3DriverConstructor = func(rootDirectory string, storageClass s3.StorageClass) (*Driver, error) {
 		encryptBool := false
 		if encrypt != "" {
 			encryptBool, err = strconv.ParseBool(encrypt)
@@ -70,7 +70,7 @@ func init() {
 			v4AuthBool,
 			minChunkSize,
 			rootDirectory,
-			s3.StandardStorage,
+			storageClass,
 			"",
 		}
 
@@ -86,7 +86,7 @@ func init() {
 	}
 
 	testsuites.RegisterSuite(func() (storagedriver.StorageDriver, error) {
-		return s3DriverConstructor(root)
+		return s3DriverConstructor(root, s3.StandardStorage)
 	}, skipS3)
 }
 
@@ -101,17 +101,17 @@ func TestEmptyRootList(t *testing.T) {
 	}
 	defer os.Remove(validRoot)
 
-	rootedDriver, err := s3DriverConstructor(validRoot)
+	rootedDriver, err := s3DriverConstructor(validRoot, s3.StandardStorage)
 	if err != nil {
 		t.Fatalf("unexpected error creating rooted driver: %v", err)
 	}
 
-	emptyRootDriver, err := s3DriverConstructor("")
+	emptyRootDriver, err := s3DriverConstructor("", s3.StandardStorage)
 	if err != nil {
 		t.Fatalf("unexpected error creating empty root driver: %v", err)
 	}
 
-	slashRootDriver, err := s3DriverConstructor("/")
+	slashRootDriver, err := s3DriverConstructor("/", s3.StandardStorage)
 	if err != nil {
 		t.Fatalf("unexpected error creating slash root driver: %v", err)
 	}
@@ -139,3 +139,63 @@ func TestEmptyRootList(t *testing.T) {
 		}
 	}
 }
+
+func TestStorageClass(t *testing.T) {
+	if skipS3() != "" {
+		t.Skip(skipS3())
+	}
+
+	rootDir, err := ioutil.TempDir("", "driver-")
+	if err != nil {
+		t.Fatalf("unexpected error creating temporary directory: %v", err)
+	}
+	defer os.Remove(rootDir)
+
+	standardDriver, err := s3DriverConstructor(rootDir, s3.StandardStorage)
+	if err != nil {
+		t.Fatalf("unexpected error creating driver with standard storage: %v", err)
+	}
+
+	rrDriver, err := s3DriverConstructor(rootDir, s3.ReducedRedundancy)
+	if err != nil {
+		t.Fatalf("unexpected error creating driver with reduced redundancy storage: %v", err)
+	}
+
+	standardFilename := "/test-standard"
+	rrFilename := "/test-rr"
+	contents := []byte("contents")
+	ctx := context.Background()
+
+	err = standardDriver.PutContent(ctx, standardFilename, contents)
+	if err != nil {
+		t.Fatalf("unexpected error creating content: %v", err)
+	}
+	defer standardDriver.Delete(ctx, standardFilename)
+
+	err = rrDriver.PutContent(ctx, rrFilename, contents)
+	if err != nil {
+		t.Fatalf("unexpected error creating content: %v", err)
+	}
+	defer rrDriver.Delete(ctx, rrFilename)
+
+	standardDriverUnwrapped := standardDriver.Base.StorageDriver.(*driver)
+	resp, err := standardDriverUnwrapped.Bucket.GetResponse(standardDriverUnwrapped.s3Path(standardFilename))
+	if err != nil {
+		t.Fatalf("unexpected error retrieving standard storage file: %v", err)
+	}
+	defer resp.Body.Close()
+	// Amazon only populates this header value for non-standard storage classes
+	if storageClass := resp.Header.Get("x-amz-storage-class"); storageClass != "" {
+		t.Fatalf("unexpected storage class for standard file: %v", storageClass)
+	}
+
+	rrDriverUnwrapped := rrDriver.Base.StorageDriver.(*driver)
+	resp, err = rrDriverUnwrapped.Bucket.GetResponse(rrDriverUnwrapped.s3Path(rrFilename))
+	if err != nil {
+		t.Fatalf("unexpected error retrieving reduced-redundancy storage file: %v", err)
+	}
+	defer resp.Body.Close()
+	if storageClass := resp.Header.Get("x-amz-storage-class"); storageClass != string(s3.ReducedRedundancy) {
+		t.Fatalf("unexpected storage class for standard file: %v", storageClass)
+	}
+}