Merge pull request #2471 from ujjwal96/traceback-fix

Fixed addon error in test suite

4 files changed, 15 insertions, 11 deletions

diff --git a/test/mitmproxy/addons/test_cut.py b/test/mitmproxy/addons/test_cut.py
index e028331f..242c6c2f 100644
--- a/test/mitmproxy/addons/test_cut.py
+++ b/test/mitmproxy/addons/test_cut.py
@@ -32,7 +32,7 @@ def test_extract():
         ["s.header[header-response]", "svalue"],
 
         ["cc.address.port", "22"],
-        ["cc.address.host", "address"],
+        ["cc.address.host", "127.0.0.1"],
         ["cc.tls_version", "TLSv1.2"],
         ["cc.sni", "address"],
         ["cc.ssl_established", "false"],
diff --git a/test/mitmproxy/test_connections.py b/test/mitmproxy/test_connections.py
index 99367bb6..83f0bd34 100644
--- a/test/mitmproxy/test_connections.py
+++ b/test/mitmproxy/test_connections.py
@@ -29,7 +29,7 @@ class TestClientConnection:
 
     def test_repr(self):
         c = tflow.tclient_conn()
-        assert 'address:22' in repr(c)
+        assert '127.0.0.1:22' in repr(c)
         assert 'ALPN' in repr(c)
         assert 'TLS' not in repr(c)
 
diff --git a/test/mitmproxy/test_flowfilter.py b/test/mitmproxy/test_flowfilter.py
index fe9b2408..c411258a 100644
--- a/test/mitmproxy/test_flowfilter.py
+++ b/test/mitmproxy/test_flowfilter.py
@@ -225,11 +225,11 @@ class TestMatchingHTTPFlow:
 
     def test_src(self):
         q = self.req()
-        assert self.q("~src address", q)
+        assert self.q("~src 127.0.0.1", q)
         assert not self.q("~src foobar", q)
         assert self.q("~src :22", q)
         assert not self.q("~src :99", q)
-        assert self.q("~src address:22", q)
+        assert self.q("~src 127.0.0.1:22", q)
 
         q.client_conn.address = None
         assert not self.q('~src address:22', q)
@@ -315,11 +315,11 @@ class TestMatchingTCPFlow:
 
     def test_src(self):
         f = self.flow()
-        assert self.q("~src address", f)
+        assert self.q("~src 127.0.0.1", f)
         assert not self.q("~src foobar", f)
         assert self.q("~src :22", f)
         assert not self.q("~src :99", f)
-        assert self.q("~src address:22", f)
+        assert self.q("~src 127.0.0.1:22", f)
 
     def test_dst(self):
         f = self.flow()
@@ -440,11 +440,11 @@ class TestMatchingWebSocketFlow:
 
     def test_src(self):
         f = self.flow()
-        assert self.q("~src address", f)
+        assert self.q("~src 127.0.0.1", f)
         assert not self.q("~src foobar", f)
         assert self.q("~src :22", f)
         assert not self.q("~src :99", f)
-        assert self.q("~src address:22", f)
+        assert self.q("~src 127.0.0.1:22", f)
 
     def test_dst(self):
         f = self.flow()
@@ -524,7 +524,7 @@ class TestMatchingDummyFlow:
 
         assert not self.q("~s", f)
 
-        assert self.q("~src address", f)
+        assert self.q("~src 127.0.0.1", f)
         assert not self.q("~src nonexistent", f)
 
         assert not self.q("~tcp", f)
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 3a2050e1..9faaf20e 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -2,6 +2,7 @@ import os.path
 import threading
 import tempfile
 import sys
+from unittest import mock
 
 import mitmproxy.platform
 from mitmproxy.proxy.config import ProxyConfig
@@ -23,10 +24,13 @@ class MasterTest:
 
     def cycle(self, master, content):
         f = tflow.tflow(req=tutils.treq(content=content))
-        master.addons.handle_lifecycle("clientconnect", f.client_conn)
+        layer = mock.Mock("mitmproxy.proxy.protocol.base.Layer")
+        layer.client_conn = f.client_conn
+        layer.reply = controller.DummyReply()
+        master.addons.handle_lifecycle("clientconnect", layer)
         for i in eventsequence.iterate(f):
             master.addons.handle_lifecycle(*i)
-        master.addons.handle_lifecycle("clientdisconnect", f.client_conn)
+        master.addons.handle_lifecycle("clientdisconnect", layer)
         return f
 
     def dummy_cycle(self, master, n, content):