From 51c8bf984b88ce60a3bd1a4859e7c80608fa26c9 Mon Sep 17 00:00:00 2001
From: Tres Seaver <tseaver@palladion.com>
Date: Fri, 6 Nov 2015 14:14:32 -0500
Subject: [PATCH] Allow setting 'bucket.storage_class'.

Fixes #1030.
---
 gcloud/storage/bucket.py      | 18 ++++++++++++++++++
 gcloud/storage/test_bucket.py | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/gcloud/storage/bucket.py b/gcloud/storage/bucket.py
index 4329e978ed0f..f29b510dff1f 100644
--- a/gcloud/storage/bucket.py
+++ b/gcloud/storage/bucket.py
@@ -89,6 +89,8 @@ class Bucket(_PropertyMixin):
     This is used in Bucket.delete() and Bucket.make_public().
     """
 
+    _STORAGE_CLASSES = ('STANDARD', 'NEARLINE', 'DURABLE_REDUCED_AVAILABILITY')
+
     def __init__(self, client, name=None):
         super(Bucket, self).__init__(name=name)
         self._client = client
@@ -675,6 +677,22 @@ def storage_class(self):
         """
         return self._properties.get('storageClass')
 
+    @storage_class.setter
+    def storage_class(self, value):
+        """Set the storage class for the bucket.
+
+        See: https://cloud.google.com/storage/docs/storage-classes
+        https://cloud.google.com/storage/docs/nearline-storage
+        https://cloud.google.com/storage/docs/durable-reduced-availability
+
+        :type value: string
+        :param value: one of "STANDARD", "NEARLINE", or
+                      "DURABLE_REDUCED_AVAILABILITY"
+        """
+        if value not in self._STORAGE_CLASSES:
+            raise ValueError('Invalid storage class: %s' % (value,))
+        self._patch_property('storageClass', value)
+
     @property
     def time_created(self):
         """Retrieve the timestamp at which the bucket was created.
diff --git a/gcloud/storage/test_bucket.py b/gcloud/storage/test_bucket.py
index 97f213be6f21..947b1f0e600d 100644
--- a/gcloud/storage/test_bucket.py
+++ b/gcloud/storage/test_bucket.py
@@ -569,6 +569,7 @@ def test_location_setter(self):
         self.assertEqual(bucket.location, None)
         bucket.location = 'AS'
         self.assertEqual(bucket.location, 'AS')
+        self.assertTrue('location' in bucket._changes)
 
     def test_lifecycle_rules_getter(self):
         NAME = 'name'
@@ -588,6 +589,7 @@ def test_lifecycle_rules_setter(self):
         self.assertEqual(bucket.lifecycle_rules, [])
         bucket.lifecycle_rules = rules
         self.assertEqual(bucket.lifecycle_rules, rules)
+        self.assertTrue('lifecycle' in bucket._changes)
 
     def test_cors_getter(self):
         NAME = 'name'
@@ -619,6 +621,7 @@ def test_cors_setter(self):
         self.assertEqual(bucket.cors, [])
         bucket.cors = [CORS_ENTRY]
         self.assertEqual(bucket.cors, [CORS_ENTRY])
+        self.assertTrue('cors' in bucket._changes)
 
     def test_get_logging_w_prefix(self):
         NAME = 'name'
@@ -712,12 +715,40 @@ def test_self_link(self):
         bucket = self._makeOne(properties=properties)
         self.assertEqual(bucket.self_link, SELF_LINK)
 
-    def test_storage_class(self):
+    def test_storage_class_getter(self):
         STORAGE_CLASS = 'http://example.com/self/'
         properties = {'storageClass': STORAGE_CLASS}
         bucket = self._makeOne(properties=properties)
         self.assertEqual(bucket.storage_class, STORAGE_CLASS)
 
+    def test_storage_class_setter_invalid(self):
+        NAME = 'name'
+        bucket = self._makeOne(name=NAME)
+        with self.assertRaises(ValueError):
+            bucket.storage_class = 'BOGUS'
+        self.assertFalse('storageClass' in bucket._changes)
+
+    def test_storage_class_setter_STANDARD(self):
+        NAME = 'name'
+        bucket = self._makeOne(name=NAME)
+        bucket.storage_class = 'STANDARD'
+        self.assertEqual(bucket.storage_class, 'STANDARD')
+        self.assertTrue('storageClass' in bucket._changes)
+
+    def test_storage_class_setter_NEARLINE(self):
+        NAME = 'name'
+        bucket = self._makeOne(name=NAME)
+        bucket.storage_class = 'NEARLINE'
+        self.assertEqual(bucket.storage_class, 'NEARLINE')
+        self.assertTrue('storageClass' in bucket._changes)
+
+    def test_storage_class_setter_DURABLE_REDUCED_AVAILABILITY(self):
+        NAME = 'name'
+        bucket = self._makeOne(name=NAME)
+        bucket.storage_class = 'DURABLE_REDUCED_AVAILABILITY'
+        self.assertEqual(bucket.storage_class, 'DURABLE_REDUCED_AVAILABILITY')
+        self.assertTrue('storageClass' in bucket._changes)
+
     def test_time_created(self):
         import datetime
         from gcloud._helpers import _RFC3339_MICROS