Merge lp:~abentley/workspace-runner/upload-mime-type into lp:workspace-runner

[Aaron Bentley]

Set content-types for uploads.

In this branch, the upload script can now use file extension to determine Content-Type header.

[Curtis Hovey (sinzui)]

Thank you. I have a question inline.

review: Needs Information (code)

[Curtis Hovey (sinzui)]

We spoke on irc. The mimetypes modules could reduce the number of lines, but we will still need to register some of the special types we have. Can you make sure that I can read *.changes and *.dsc in my browser?

review: Approve (code)

=== modified file 'workspace_runner/tests/test_upload_artifacts.py'
--- workspace_runner/tests/test_upload_artifacts.py	2015-06-29 19:19:16 +0000
+++ workspace_runner/tests/test_upload_artifacts.py	2015-06-30 20:28:44 +0000
@@ -1,4 +1,5 @@
 from argparse import Namespace
+import errno
 import json
 import os
 from tempfile import NamedTemporaryFile
@@ -12,6 +13,7 @@
     parse_args,
     resolve_rules,
     upload_artifacts,
+    upload_file,
     )
 
 
@@ -75,9 +77,42 @@
     def __init__(self, path):
         self.contents = None
         self.path = path
+        self.headers = None
 
-    def set_contents_from_file(self, fp):
+    def set_contents_from_file(self, fp, headers=None):
         self.contents = fp.read()
+        self.headers = headers
+
+
+class TestUploadFile(TestCase):
+
+    def ct_test(self, mime_type, extension=None, filename=None):
+        bucket = FakeBucket()
+        with temp_dir() as files_dir:
+            if filename is None:
+                filename = 'foo.{}'.format(extension)
+            file_path = os.path.join(files_dir, filename)
+            parent = os.path.dirname(file_path)
+            try:
+                os.makedirs(parent)
+            except OSError as e:
+                if e.errno != errno.EEXIST:
+                    raise
+            open(file_path, 'w').close()
+            upload_file(file_path, 'foo', 'bar', bucket)
+        self.assertEqual(bucket.keys[0].headers,
+                         {'Content-Type': mime_type})
+
+    def test_content_type(self):
+        self.ct_test('application/gzip', 'tar.gz')
+        self.ct_test('application/gzip', 'tgz')
+        self.ct_test('application/json', 'json')
+        self.ct_test('text/plain', 'txt')
+        self.ct_test('application/octet-stream', 'foo')
+        self.ct_test('text/x-yaml; charset=utf-8', 'jenv')
+        self.ct_test('application/zip', 'zip')
+        self.ct_test('text/plain', 'log')
+        self.ct_test('text/plain', filename='foo/consoleText')
 
 
 class TestUploadArtifacts(TestCase):
@@ -109,6 +144,7 @@
         key = bucket.keys[0]
         self.assertEqual(key.path, 'prefix1/text-files/foo.txt')
         self.assertEqual(key.contents, 'foo file')
+        self.assertEqual(key.headers, {'Content-Type': 'text/plain'})
 
 
 class TestMain(TestCase):
 
=== modified file 'workspace_runner/upload_artifacts.py'
--- workspace_runner/upload_artifacts.py	2015-06-30 18:37:21 +0000
+++ workspace_runner/upload_artifacts.py	2015-06-30 20:28:44 +0000
@@ -27,13 +27,39 @@
     return result
 
 
+def get_media_type(source_name):
+    """Determine a media type from a filename."""
+    basename = os.path.basename(source_name)
+    if basename == 'consoleText':
+        return 'text/plain'
+    split_name = basename.rsplit('.', 1)
+    if len(split_name) == 1:
+        return 'application/octet-stream'
+    else:
+        extension = split_name[-1]
+    return {
+        'gz': 'application/gzip',
+        'jenv': 'text/x-yaml; charset=utf-8',
+        'json': 'application/json',
+        'log': 'text/plain',
+        'tgz': 'application/gzip',
+        'txt': 'text/plain',
+        'zip': 'application/zip',
+    }.get(extension, 'application/octet-stream')
+
+
+def upload_file(source_name, prefix, dest, bucket):
+    media_type = get_media_type(source_name)
+    key = bucket.new_key('/'.join((prefix, dest)))
+    with open(source_name) as source:
+        key.set_contents_from_file(source, {'Content-Type': media_type})
+
+
 def upload_artifacts(root, settings, s3_factory=S3Connection):
     connection = s3_factory(settings['access_key'], settings['secret_key'])
     bucket = connection.get_bucket(settings['bucket'])
     for source_name, dest in resolve_rules(root, settings['files']):
-        key = bucket.new_key('/'.join((settings['prefix'], dest)))
-        with open(source_name) as source:
-            key.set_contents_from_file(source)
+        upload_file(source_name, settings['prefix'], dest, bucket)
 
 
 def main(argv=None):