support multiple boto versions for header writing

Depending on the version of boto, different methods are necessary for
manipulating headers. We now support 2.0rc1 and 2.0 by sub-classing
boto.s3.connection.S3Connection and overloading
boto.connection.HTTPRequest.authorize, respectfully.

s3tests/functional/test_headers.py

@@ -1,4 +1,5 @@
 from cStringIO import StringIO
+import boto.connection
 import boto.exception
 import boto.s3.connection
 import boto.s3.acl
@@ -31,12 +32,15 @@ from . import (
 
 
 _orig_conn = {}
+_orig_authorize = None
 _custom_headers = {}
 _remove_headers = []
+boto_type = None
 
 
-# fill_in_auth does not exist in boto master, so if we update boto, this will
-# need to be changed to handle make_request's changes. Likely, at authorize
+# HeaderS3Connection and _our_authorize are necessary to be able to arbitrarily
+# overwrite headers. Depending on the version of boto, one or the other is
+# necessary. We later determine in setup what needs to be used.
 class HeaderS3Connection(S3Connection):
     def fill_in_auth(self, http_request, **kwargs):
         global _custom_headers, _remove_headers
@@ -65,9 +69,41 @@ class HeaderS3Connection(S3Connection):
         return http_request
 
 
+def _our_authorize(self, connection, **kwargs):
+    global _custom_headers, _remove_headers
+
+    # do our header magic
+    final_headers = self.headers
+    final_headers.update(_custom_headers)
+
+    for header in _remove_headers:
+        try:
+            del final_headers[header]
+        except KeyError:
+            pass
+
+    _orig_authorize(self, connection, **kwargs)
+
+    final_headers = self.headers
+    final_headers.update(_custom_headers)
+
+    for header in _remove_headers:
+        try:
+            del final_headers[header]
+        except KeyError:
+            pass
+
+
 def setup():
+    global boto_type
+
+    # we determine what we need to replace by the existence of particular
+    # attributes. boto 2.0rc1 as fill_in_auth for S3Connection, while boto 2.0
+    # has authorize for HTTPRequest.
+    if hasattr(S3Connection, 'fill_in_auth'):
         global _orig_conn
 
+        boto_type = 'S3Connection'
         for conn in s3:
             _orig_conn[conn] = s3[conn]
             header_conn = HeaderS3Connection(
@@ -80,12 +116,33 @@ def setup():
                 )
 
             s3[conn] = header_conn
+    elif hasattr(boto.connection.HTTPRequest, 'authorize'):
+        global _orig_authorize
+
+        boto_type = 'HTTPRequest'
+
+        _orig_authorize = boto.connection.HTTPRequest.authorize
+        boto.connection.HTTPRequest.authorize = _our_authorize
+    else:
+        raise RuntimeError
 
 
 def teardown():
-    global _orig_auth_handler
+    global boto_type
+
+    # replace original functionality depending on the boto version
+    if boto_type is 'S3Connection':
+        global _orig_conn
         for conn in s3:
             s3[conn] = _orig_conn[conn]
+        _orig_conn = {}
+    elif boto_type is 'HTTPRequest':
+        global _orig_authorize
+
+        boto.connection.HTTPRequest.authorize = _orig_authorize
+        _orig_authorize = None
+    else:
+        raise RuntimeError
 
 
 def _clear_custom_headers():