From 3c305bbe35a0dd63bacc3e35aa0a5ed548f90c91 Mon Sep 17 00:00:00 2001
From: Eddie Hebert <ehebert@quantopian.com>
Date: Mon, 29 Jul 2013 14:30:42 -0400
Subject: [PATCH] TST: Add script to automate upload of risk answer sheet to
 S3.

For maintainer use, requires AWS credentials for the account where
the `zipline-test-data` bucket is hosted.

Script does the following steps which used to be manual:
- Create a key name based on the md5 of the answer key file.
- Upload the answer key to S3 bucket.
- Make the file publically downloadable over HTTP.
---
 tests/risk/answer_key.py        | 26 +++++++++------
 tests/risk/upload_answer_key.py | 57 +++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 9 deletions(-)
 create mode 100644 tests/risk/upload_answer_key.py

diff --git a/tests/risk/answer_key.py b/tests/risk/answer_key.py
index 8583bcba..b832e131 100644
--- a/tests/risk/answer_key.py
+++ b/tests/risk/answer_key.py
@@ -34,13 +34,28 @@ DIR = os.path.dirname(os.path.realpath(__file__))
 ANSWER_KEY_CHECKSUMS_PATH = os.path.join(DIR, 'risk-answer-key-checksums')
 ANSWER_KEY_CHECKSUMS = open(ANSWER_KEY_CHECKSUMS_PATH, 'r').read().splitlines()
 
-ANSWER_KEY_PATH = os.path.join(DIR, 'risk-answer-key.xlsx')
+ANSWER_KEY_FILENAME = 'risk-answer-key.xlsx'
+
+ANSWER_KEY_PATH = os.path.join(DIR, ANSWER_KEY_FILENAME)
+
+ANSWER_KEY_BUCKET_NAME = 'zipline-test_data'
 
 ANSWER_KEY_DL_TEMPLATE = """
 https://s3.amazonaws.com/zipline-test-data/risk/{md5}/risk-answer-key.xlsx
 """.strip()
 
 
+def answer_key_signature():
+    with open(ANSWER_KEY_PATH, 'r') as f:
+        md5 = hashlib.md5()
+        while True:
+            buf = f.read(1024)
+            if not buf:
+                break
+            md5.update(buf)
+    return md5.hexdigest()
+
+
 def ensure_latest_answer_key():
     """
     Get the latest answer key from a publically available location.
@@ -70,14 +85,7 @@ def ensure_latest_answer_key():
 
     local_answer_key_exists = os.path.exists(ANSWER_KEY_PATH)
     if local_answer_key_exists:
-        with open(ANSWER_KEY_PATH, 'r') as f:
-            md5 = hashlib.md5()
-            while True:
-                buf = f.read(1024)
-                if not buf:
-                    break
-                md5.update(buf)
-        local_hash = md5.hexdigest()
+        local_hash = answer_key_signature()
 
         if local_hash in ANSWER_KEY_CHECKSUMS:
             # Assume previously downloaded version.
diff --git a/tests/risk/upload_answer_key.py b/tests/risk/upload_answer_key.py
new file mode 100644
index 00000000..fadab1bb
--- /dev/null
+++ b/tests/risk/upload_answer_key.py
@@ -0,0 +1,57 @@
+#
+# Copyright 2013 Quantopian, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Utility script for maintainer use to upload current version of the answer key
+spreadsheet to S3.
+"""
+import hashlib
+
+import boto
+
+from . import answer_key
+
+BUCKET_NAME = 'zipline-test-data'
+
+
+def main():
+    with open(answer_key.ANSWER_KEY_PATH, 'r') as f:
+        md5 = hashlib.md5()
+        while True:
+            buf = f.read(1024)
+            if not buf:
+                break
+            md5.update(buf)
+    local_hash = md5.hexdigest()
+
+    s3_conn = boto.connect_s3()
+
+    bucket = s3_conn.get_bucket(BUCKET_NAME)
+    key = boto.s3.key.Key(bucket)
+
+    key.key = "risk/{local_hash}/risk-answer-key.xlsx".format(
+        local_hash=local_hash)
+    key.set_contents_from_filename(answer_key.ANSWER_KEY_PATH)
+    key.set_acl('public-read')
+
+    download_link = "http://s3.amazonaws.com/{bucket_name}/{key}".format(
+        bucket_name=BUCKET_NAME,
+        key=key.key)
+
+    print("Uploaded to key: {key}".format(key=key.key))
+    print("Download link: {download_link}".format(download_link=download_link))
+
+if __name__ == "__main__":
+    main()