Merge pull request #1511 from arjun23496/count_in_replace

Fixes #1495 - Added count argument for replacing contents in body

3 files changed, 25 insertions, 0 deletions

diff --git a/test/netlib/http/test_headers.py b/test/netlib/http/test_headers.py
index 51537310..ad2bc548 100644
--- a/test/netlib/http/test_headers.py
+++ b/test/netlib/http/test_headers.py
@@ -75,6 +75,11 @@ 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..74272309 100644
--- a/test/netlib/http/test_message.py
+++ b/test/netlib/http/test_message.py
@@ -99,6 +99,16 @@ 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..9baabaa6 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")