Merge pull request #2120 from ujjwal96/replace-from-file

Replace from file through replacement editor #1458

1 files changed, 36 insertions, 33 deletions

diff --git a/test/mitmproxy/addons/test_replace.py b/test/mitmproxy/addons/test_replace.py
index 2311641a..7d590b35 100644
--- a/test/mitmproxy/addons/test_replace.py
+++ b/test/mitmproxy/addons/test_replace.py
@@ -1,6 +1,5 @@
 import pytest
 
-from .. import tservers
 from mitmproxy.addons import replace
 from mitmproxy.test import taddons
 from mitmproxy.test import tflow
@@ -32,7 +31,7 @@ class TestReplace:
         with taddons.context() as tctx:
             tctx.configure(
                 r,
-                replacements = [
+                replacements=[
                     "/~q/foo/bar",
                     "/~s/foo/bar",
                 ]
@@ -47,53 +46,57 @@ class TestReplace:
             r.response(f)
             assert f.response.content == b"bar"
 
-
-class TestUpstreamProxy(tservers.HTTPUpstreamProxyTest):
-    ssl = False
-
     def test_order(self):
-        sa = replace.Replace()
-        self.proxy.tmaster.addons.add(sa)
-
-        self.proxy.tmaster.options.replacements = [
-            "/~q/foo/bar",
-            "/~q/bar/baz",
-            "/~q/foo/oh noes!",
-            "/~s/baz/ORLY"
-        ]
-        p = self.pathoc()
-        with p.connect():
-            req = p.request("get:'%s/p/418:b\"foo\"'" % self.server.urlbase)
-        assert req.content == b"ORLY"
-        assert req.status_code == 418
+        r = replace.Replace()
+        with taddons.context() as tctx:
+            tctx.configure(
+                r,
+                replacements=[
+                    "/foo/bar",
+                    "/bar/baz",
+                    "/foo/oh noes!",
+                    "/bar/oh noes!",
+                ]
+            )
+            f = tflow.tflow()
+            f.request.content = b"foo"
+            r.request(f)
+            assert f.request.content == b"baz"
 
 
 class TestReplaceFile:
     def test_simple(self, tmpdir):
-        r = replace.ReplaceFile()
-        rp = tmpdir.join("replacement")
-        rp.write("bar")
+        r = replace.Replace()
         with taddons.context() as tctx:
+            tmpfile = tmpdir.join("replacement")
+            tmpfile.write("bar")
             tctx.configure(
                 r,
-                replacement_files = [
-                    "/~q/foo/" + str(rp),
-                    "/~s/foo/" + str(rp),
-                    "/~b nonexistent/nonexistent/nonexistent",
-                ]
+                replacements=["/~q/foo/@" + str(tmpfile)]
             )
             f = tflow.tflow()
             f.request.content = b"foo"
             r.request(f)
             assert f.request.content == b"bar"
 
-            f = tflow.tflow(resp=True)
-            f.response.content = b"foo"
-            r.response(f)
-            assert f.response.content == b"bar"
+    def test_nonexistent(self, tmpdir):
+        r = replace.Replace()
+        with taddons.context() as tctx:
+            with pytest.raises(Exception, match="Invalid file path"):
+                tctx.configure(
+                    r,
+                    replacements=["/~q/foo/@nonexistent"]
+                )
 
+            tmpfile = tmpdir.join("replacement")
+            tmpfile.write("bar")
+            tctx.configure(
+                r,
+                replacements=["/~q/foo/@" + str(tmpfile)]
+            )
+            tmpfile.remove()
             f = tflow.tflow()
-            f.request.content = b"nonexistent"
+            f.request.content = b"foo"
             assert not tctx.master.event_log
             r.request(f)
             assert tctx.master.event_log