tests++

2 files changed, 32 insertions, 2 deletions

diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index b9c6a2f6..5c1e6a68 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -680,6 +680,10 @@ class TestSerialize:
         for i in range(3):
             f = tutils.tflow(err=True)
             w.add(f)
+        f = tutils.ttcpflow()
+        w.add(f)
+        f = tutils.ttcpflow(err=True)
+        w.add(f)
 
         sio.seek(0)
         return flow.FlowReader(sio)
@@ -1205,6 +1209,10 @@ class TestError:
         e3 = e.copy()
         assert e3.get_state() == e.get_state()
 
+    def test_repr(self):
+        e = Error("yay")
+        assert repr(e)
+
 
 class TestClientConnection:
 
diff --git a/test/mitmproxy/tutils.py b/test/mitmproxy/tutils.py
index d51ac185..270a3911 100644
--- a/test/mitmproxy/tutils.py
+++ b/test/mitmproxy/tutils.py
@@ -3,6 +3,8 @@ import shutil
 import tempfile
 import argparse
 import sys
+
+from mitmproxy.models.tcp import TCPMessage
 from six.moves import cStringIO as StringIO
 from contextlib import contextmanager
 
@@ -11,7 +13,7 @@ from unittest.case import SkipTest
 import netlib.tutils
 from mitmproxy import utils, controller
 from mitmproxy.models import (
-    ClientConnection, ServerConnection, Error, HTTPRequest, HTTPResponse, HTTPFlow
+    ClientConnection, ServerConnection, Error, HTTPRequest, HTTPResponse, HTTPFlow, TCPFlow
 )
 
 
@@ -44,6 +46,26 @@ def skip_appveyor(fn):
         return fn
 
 
+def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None):
+    if client_conn is True:
+        client_conn = tclient_conn()
+    if server_conn is True:
+        server_conn = tserver_conn()
+    if messages is True:
+        messages = [
+            TCPMessage(True, b"hello"),
+            TCPMessage(False, b"it's me"),
+        ]
+    if err is True:
+        err = terr()
+
+    f = TCPFlow(client_conn, server_conn)
+    f.messages = messages
+    f.error = err
+    f.reply = controller.DummyReply()
+    return f
+
+
 def tflow(client_conn=True, server_conn=True, req=True, resp=None, err=None):
     """
     @type client_conn: bool | None | mitmproxy.proxy.connection.ClientConnection
@@ -51,7 +73,7 @@ def tflow(client_conn=True, server_conn=True, req=True, resp=None, err=None):
     @type req:         bool | None | mitmproxy.protocol.http.HTTPRequest
     @type resp:        bool | None | mitmproxy.protocol.http.HTTPResponse
     @type err:         bool | None | mitmproxy.protocol.primitives.Error
-    @return:           bool | None | mitmproxy.protocol.http.HTTPFlow
+    @return:           mitmproxy.protocol.http.HTTPFlow
     """
     if client_conn is True:
         client_conn = tclient_conn()