3 files changed, 10 insertions, 10 deletions
diff --git a/test/mitmproxy/scripts/tcp_stream_modify.py b/test/mitmproxy/scripts/tcp_stream_modify.py
index 93b0d5c8..d7953ef9 100644
--- a/test/mitmproxy/scripts/tcp_stream_modify.py
+++ b/test/mitmproxy/scripts/tcp_stream_modify.py
@@ -1,3 +1,4 @@
-def tcp_message(ctx, tm):
-    if tm.sender == tm.server_conn:
-        tm.message = tm.message.replace("foo", "bar")
+def tcp_message(ctx, flow):
+    message = flow.messages[-1]
+    if not message.from_client:
+        message.content = message.content.replace("foo", "bar")
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index 454736d4..0701d52b 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -14,7 +14,7 @@ from pathod import pathoc, pathod
 
 from mitmproxy.proxy.config import HostMatcher
 from mitmproxy.exceptions import Kill
-from mitmproxy.models import Error, HTTPResponse
+from mitmproxy.models import Error, HTTPResponse, HTTPFlow
 
 from . import tutils, tservers
 
@@ -177,9 +177,9 @@ class TcpMixin:
         assert n.status_code == 304
         assert i.status_code == 305
         assert i2.status_code == 306
-        assert any(f.response.status_code == 304 for f in self.master.state.flows)
-        assert not any(f.response.status_code == 305 for f in self.master.state.flows)
-        assert not any(f.response.status_code == 306 for f in self.master.state.flows)
+        assert any(f.response.status_code == 304 for f in self.master.state.flows if isinstance(f, HTTPFlow))
+        assert not any(f.response.status_code == 305 for f in self.master.state.flows if isinstance(f, HTTPFlow))
+        assert not any(f.response.status_code == 306 for f in self.master.state.flows if isinstance(f, HTTPFlow))
 
         # Test that we get the original SSL cert
         if self.ssl:
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 4fa519cc..c9d68cfd 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -50,9 +50,8 @@ class TestMaster(flow.FlowMaster):
     def clear_log(self):
         self.log = []
 
-    def handle_log(self, l):
-        self.log.append(l.msg)
-        l.reply()
+    def add_event(self, message, level=None):
+        self.log.append(message)
 
 
 class ProxyThread(threading.Thread):