This commit is contained in:
Maximilian Hils 2016-05-20 13:27:26 -07:00
parent f96697646c
commit a3946d2a2d
2 changed files with 32 additions and 2 deletions

View File

@ -680,6 +680,10 @@ class TestSerialize:
for i in range(3): for i in range(3):
f = tutils.tflow(err=True) f = tutils.tflow(err=True)
w.add(f) w.add(f)
f = tutils.ttcpflow()
w.add(f)
f = tutils.ttcpflow(err=True)
w.add(f)
sio.seek(0) sio.seek(0)
return flow.FlowReader(sio) return flow.FlowReader(sio)
@ -1205,6 +1209,10 @@ class TestError:
e3 = e.copy() e3 = e.copy()
assert e3.get_state() == e.get_state() assert e3.get_state() == e.get_state()
def test_repr(self):
e = Error("yay")
assert repr(e)
class TestClientConnection: class TestClientConnection:

View File

@ -3,6 +3,8 @@ import shutil
import tempfile import tempfile
import argparse import argparse
import sys import sys
from mitmproxy.models.tcp import TCPMessage
from six.moves import cStringIO as StringIO from six.moves import cStringIO as StringIO
from contextlib import contextmanager from contextlib import contextmanager
@ -11,7 +13,7 @@ from unittest.case import SkipTest
import netlib.tutils import netlib.tutils
from mitmproxy import utils, controller from mitmproxy import utils, controller
from mitmproxy.models import ( 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 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): def tflow(client_conn=True, server_conn=True, req=True, resp=None, err=None):
""" """
@type client_conn: bool | None | mitmproxy.proxy.connection.ClientConnection @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 req: bool | None | mitmproxy.protocol.http.HTTPRequest
@type resp: bool | None | mitmproxy.protocol.http.HTTPResponse @type resp: bool | None | mitmproxy.protocol.http.HTTPResponse
@type err: bool | None | mitmproxy.protocol.primitives.Error @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: if client_conn is True:
client_conn = tclient_conn() client_conn = tclient_conn()