improve .replace() and move it into netlib

3 files changed, 23 insertions, 4 deletions

diff --git a/test/mitmproxy/test_utils.py b/test/mitmproxy/test_utils.py
index ae6369ae..db7dec4a 100644
--- a/test/mitmproxy/test_utils.py
+++ b/test/mitmproxy/test_utils.py
@@ -99,7 +99,3 @@ def test_parse_size():
     assert utils.parse_size("1g") == 1024**3
     tutils.raises(ValueError, utils.parse_size, "1f")
     tutils.raises(ValueError, utils.parse_size, "ak")
-
-
-def test_safe_subn():
-    assert utils.safe_subn("foo", u"bar", "\xc2foo")
diff --git a/test/netlib/http/test_headers.py b/test/netlib/http/test_headers.py
index d50fee3e..8c1db9dc 100644
--- a/test/netlib/http/test_headers.py
+++ b/test/netlib/http/test_headers.py
@@ -150,3 +150,22 @@ class TestHeaders(object):
         assert headers != headers2
         headers2.set_state(headers.get_state())
         assert headers == headers2
+
+    def test_replace_simple(self):
+        headers = Headers(Host="example.com", Accept="text/plain")
+        replacements = headers.replace("Host: ", "X-Host: ")
+        assert replacements == 1
+        assert headers["X-Host"] == "example.com"
+        assert "Host" not in headers
+        assert headers["Accept"] == "text/plain"
+
+    def test_replace_multi(self):
+        headers = self._2host()
+        headers.replace(r"Host: example\.com", r"Host: example.de")
+        assert headers.get_all("Host") == ["example.de", "example.org"]
+
+    def test_replace_remove_spacer(self):
+        headers = Headers(Host="example.com")
+        replacements = headers.replace(r"Host: ", "X-Host ")
+        assert replacements == 0
+        assert headers["Host"] == "example.com"
diff --git a/test/netlib/test_utils.py b/test/netlib/test_utils.py
index fcb63eb2..be2a59fc 100644
--- a/test/netlib/test_utils.py
+++ b/test/netlib/test_utils.py
@@ -166,3 +166,7 @@ class TestSerializable:
         a.set_state(1)
         assert a.i == 1
         assert b.i == 42
+
+
+def test_safe_subn():
+    assert utils.safe_subn("foo", u"bar", "\xc2foo")