Added tests for replace in request,message,header

4 files changed, 25 insertions, 1 deletions

diff --git a/netlib/http/headers.py b/netlib/http/headers.py
index 60a36515..f4cf8b7b 100644
--- a/netlib/http/headers.py
+++ b/netlib/http/headers.py
@@ -185,6 +185,8 @@ class Headers(multidict.MultiDict):
             else:
                 count -= n
                 replacements += n
+                if count == 0:
+                    break;
             fields.append((name, value))
         self.fields = tuple(fields)
         return replacements
diff --git a/test/netlib/http/test_headers.py b/test/netlib/http/test_headers.py
index 51537310..b1085e5a 100644
--- a/test/netlib/http/test_headers.py
+++ b/test/netlib/http/test_headers.py
@@ -75,6 +75,10 @@ class TestHeaders(object):
         assert replacements == 0
         assert headers["Host"] == "example.com"
 
+    def test_replace_with_count(self):
+        headers = Headers(Host="foobarfoo.com", Accept="foo/bar")
+        replacements = headers.replace("foo","bar",count=1)
+        assert replacements == 1
 
 def test_parse_content_type():
     p = parse_content_type
diff --git a/test/netlib/http/test_message.py b/test/netlib/http/test_message.py
index 12e4706c..df66ef6e 100644
--- a/test/netlib/http/test_message.py
+++ b/test/netlib/http/test_message.py
@@ -99,6 +99,15 @@ class TestMessage(object):
     def test_http_version(self):
         _test_decoded_attr(tresp(), "http_version")
 
+    def test_replace(self):
+        r = tresp()
+        r.content = b"foofootoo"
+        r.replace(b"foo","gg")
+        assert r.content == b"ggggtoo"        
+
+        r.content = b"foofootoo"
+        r.replace(b"foo","gg",count=1)
+        assert r.content == b"ggfootoo"
 
 class TestMessageContentEncoding(object):
     def test_simple(self):
diff --git a/test/netlib/http/test_request.py b/test/netlib/http/test_request.py
index f3cd8b71..1fc5dc24 100644
--- a/test/netlib/http/test_request.py
+++ b/test/netlib/http/test_request.py
@@ -26,6 +26,16 @@ class TestRequestCore(object):
         request.host = None
         assert repr(request) == "Request(GET /path)"
 
+    def replace(self):
+        r = treq()
+        r.path = b"foobarfoo"
+        r.replace(b"foo","bar")
+        assert r.path == b"barbarbar"
+
+        r.path = b"foobarfoo"
+        r.replace(b"foo","bar",count=1)
+        assert r.path == b"barbarfoo"
+
     def test_first_line_format(self):
         _test_passthrough_attr(treq(), "first_line_format")
 
@@ -87,7 +97,6 @@ class TestRequestCore(object):
         request.host = "example.org"
         assert request.headers["Host"] == "example.org"
 
-
 class TestRequestUtils(object):
     """
     Tests for additional convenience methods.