From ddbf0464f46afe34a238566e84f6d5a3f7c559fd Mon Sep 17 00:00:00 2001
From: Aron Griffis <aron@arongriffis.com>
Date: Tue, 25 Aug 2015 07:12:08 -0400
Subject: [PATCH] Add replace_headers()

---
 tests/unit/test_filters.py | 61 +++++++++++++++++++++++++++++++++-----
 vcr/filters.py             | 27 ++++++++++++++---
 2 files changed, 76 insertions(+), 12 deletions(-)

diff --git a/tests/unit/test_filters.py b/tests/unit/test_filters.py
index 30bb203..c426e7c 100644
--- a/tests/unit/test_filters.py
+++ b/tests/unit/test_filters.py
@@ -1,26 +1,71 @@
 from vcr.filters import (
-    remove_headers,
+    remove_headers, replace_headers,
     remove_query_parameters,
     remove_post_data_parameters
 )
+from vcr.compat import mock
 from vcr.request import Request
 import json
 
 
+def test_replace_headers():
+    # This tests all of:
+    #   1. keeping a header
+    #   2. removing a header
+    #   3. replacing a header
+    #   4. replacing a header using a callable
+    #   5. removing a header using a callable
+    #   6. replacing a header that doesn't exist
+    headers = {
+        'one': ['keep'],
+        'two': ['lose'],
+        'three': ['change'],
+        'four': ['shout'],
+        'five': ['whisper'],
+    }
+    request = Request('GET', 'http://google.com', '', headers)
+    replace_headers(request, [
+        ('two', None),
+        ('three', 'tada'),
+        ('four', lambda key, value, request: value.upper()),
+        ('five', lambda key, value, request: None),
+        ('six', 'doesntexist'),
+    ])
+    assert request.headers == {
+        'one': 'keep',
+        'three': 'tada',
+        'four': 'SHOUT',
+    }
+
+
+def test_replace_headers_empty():
+    headers = {'hello': 'goodbye', 'secret': 'header'}
+    request = Request('GET', 'http://google.com', '', headers)
+    replace_headers(request, [])
+    assert request.headers == headers
+
+
+def test_replace_headers_callable():
+    # This goes beyond test_replace_headers() to ensure that the callable
+    # receives the expected arguments.
+    headers = {'hey': 'there'}
+    request = Request('GET', 'http://google.com', '', headers)
+    callme = mock.Mock(return_value='ho')
+    replace_headers(request, [('hey', callme)])
+    assert request.headers == {'hey': 'ho'}
+    assert callme.call_args == ((), {'request': request,
+                                     'key': 'hey',
+                                     'value': 'there'})
+
+
 def test_remove_headers():
+    # Test the backward-compatible API wrapper.
     headers = {'hello': ['goodbye'], 'secret': ['header']}
     request = Request('GET', 'http://google.com', '', headers)
     remove_headers(request, ['secret'])
     assert request.headers == {'hello': 'goodbye'}
 
 
-def test_remove_headers_empty():
-    headers = {'hello': 'goodbye', 'secret': 'header'}
-    request = Request('GET', 'http://google.com', '', headers)
-    remove_headers(request, [])
-    assert request.headers == headers
-
-
 def test_remove_query_parameters():
     uri = 'http://g.com/?q=cowboys&w=1'
     request = Request('GET', uri, '', {})
diff --git a/vcr/filters.py b/vcr/filters.py
index 10830c6..6433ddb 100644
--- a/vcr/filters.py
+++ b/vcr/filters.py
@@ -4,16 +4,35 @@ import json
 
 from .compat import collections
 
-
-def remove_headers(request, headers_to_remove):
+def replace_headers(request, replacements):
+    """
+    Replace headers in request according to replacements. The replacements
+    should be a list of (key, value) pairs where the value can be any of:
+      1. A simple replacement string value.
+      2. None to remove the given header.
+      3. A callable which accepts (key, value, request) and returns a string
+         value or None.
+    """
     new_headers = request.headers.copy()
-    for k in headers_to_remove:
+    for k, rv in replacements:
         if k in new_headers:
-            del new_headers[k]
+            ov = new_headers.pop(k)
+            if callable(rv):
+                rv = rv(key=k, value=ov, request=request)
+            if rv is not None:
+                new_headers[k] = rv
     request.headers = new_headers
     return request
 
 
+def remove_headers(request, headers_to_remove):
+    """
+    Wrap replace_headers() for API backward compatibility.
+    """
+    replacements = [(k, None) for k in headers_to_remove]
+    return replace_headers(request, replacements)
+
+
 def remove_query_parameters(request, query_parameters_to_remove):
     query = request.query
     new_query = [(k, v) for (k, v) in query