diff options
-rw-r--r-- | test/mitmproxy/test_flow.py | 8 | ||||
-rw-r--r-- | test/mitmproxy/tutils.py | 26 |
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() |