lint, mypy, tests++

7 files changed, 76 insertions, 23 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index 4b7e076b..4d0a7ef9 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -575,7 +575,7 @@ class View(collections.abc.Sequence):
     def kill(self, f):
         self.update([f])
 
-    def tcp_start(self,f):
+    def tcp_start(self, f):
         self.add([f])
 
     def tcp_message(self, f):
diff --git a/mitmproxy/tools/console/common.py b/mitmproxy/tools/console/common.py
index ccfda186..cba3a355 100644
--- a/mitmproxy/tools/console/common.py
+++ b/mitmproxy/tools/console/common.py
@@ -1,8 +1,6 @@
 import enum
 import platform
 import typing
-import datetime
-import time
 import math
 from functools import lru_cache
 from publicsuffix2 import get_sld, get_tld
@@ -343,7 +341,7 @@ def format_left_indicators(
         intercepted: bool,
         timestamp: float
 ):
-    indicators = []
+    indicators: typing.List[typing.Union[str, typing.Tuple[str, str]]] = []
     if focused:
         indicators.append(("focus", ">>"))
     else:
@@ -361,7 +359,7 @@ def format_right_indicators(
         replay: bool,
         marked: bool
 ):
-    indicators = []
+    indicators: typing.List[typing.Union[str, typing.Tuple[str, str]]] = []
     if replay:
         indicators.append(("replay", SYMBOL_REPLAY))
     else:
@@ -440,14 +438,14 @@ def format_http_flow_list(
         if intercepted:
             style = "intercept"
         else:
-            style = None
+            style = ""
 
         status_style = style or HTTP_RESPONSE_CODE_STYLE.get(response_code // 100, "code_other")
         resp.append(fcol(SYMBOL_RETURN, status_style))
         if response_is_replay:
             resp.append(fcol(SYMBOL_REPLAY, "replay"))
         resp.append(fcol(str(response_code), status_style))
-        if render_mode is RenderMode.DETAILVIEW:
+        if response_reason and render_mode is RenderMode.DETAILVIEW:
             resp.append(fcol(response_reason, status_style))
 
         if response_content_type:
@@ -512,7 +510,7 @@ def format_http_flow_table(
     if intercepted and not response_code:
         request_style = "intercept"
     else:
-        request_style = None
+        request_style = ""
 
     scheme_style = request_style or SCHEME_STYLES.get(request_scheme, "scheme_other")
     items.append(fcol(fixlen(request_scheme.upper(), 5), scheme_style))
@@ -529,7 +527,7 @@ def format_http_flow_table(
     if intercepted and response_code:
         response_style = "intercept"
     else:
-        response_style = None
+        response_style = ""
 
     if response_code:
 
@@ -584,16 +582,17 @@ def format_http_flow_table(
 
 
 @lru_cache(maxsize=800)
-def raw_format_tcp_flow(
+def format_tcp_flow(
         *,
         render_mode: RenderMode,
-        focused: typing.Optional[bool],
+        focused: bool,
         timestamp_start: float,
         marked: bool,
         client_address,
         server_address,
         total_size: int,
         duration: typing.Optional[float],
+        error_message: typing.Optional[str],
 ):
     conn = f"{human.format_address(client_address)} <-> {human.format_address(server_address)}"
 
@@ -615,6 +614,8 @@ def raw_format_tcp_flow(
         items.append(fcol("TCP", SCHEME_STYLES["tcp"]))
 
     items.append(('weight', 1.0, truncated_plain(conn, "text", 'left')))
+    if error_message:
+        items.append(('weight', 1.0, truncated_plain(error_message, "error", 'left')))
 
     if total_size:
         size, size_style = format_size(total_size)
@@ -648,6 +649,12 @@ def format_flow(
     relevant for display and call the render with only that. This assures that rows
     are updated if the flow is changed.
     """
+    duration: typing.Optional[float]
+    error_message: typing.Optional[str]
+    if f.error:
+        error_message = f.error.msg
+    else:
+        error_message = None
 
     if isinstance(f, TCPFlow):
         total_size = 0
@@ -657,7 +664,7 @@ def format_flow(
             duration = f.messages[-1].timestamp - f.timestamp_start
         else:
             duration = None
-        return raw_format_tcp_flow(
+        return format_tcp_flow(
             render_mode=render_mode,
             focused=focused,
             timestamp_start=f.timestamp_start,
@@ -666,11 +673,13 @@ def format_flow(
             server_address=f.server_conn.address,
             total_size=total_size,
             duration=duration,
+            error_message=error_message,
         )
     elif isinstance(f, HTTPFlow):
         intercepted = (
                 f.intercepted and not (f.reply and f.reply.state == "committed")
         )
+        response_content_length: typing.Optional[int]
         if f.response:
             if f.response.raw_content is not None:
                 response_content_length = len(f.response.raw_content)
@@ -691,10 +700,6 @@ def format_flow(
             response_content_type = None
             response_is_replay = False
             duration = None
-        if f.error:
-            error_message = f.error.msg
-        else:
-            error_message = None
 
         if render_mode in (RenderMode.LIST, RenderMode.DETAILVIEW):
             render_func = format_http_flow_list
@@ -720,7 +725,7 @@ def format_flow(
             response_content_type=response_content_type,
             response_is_replay=response_is_replay,
             duration=duration,
-            error_message=error_message
+            error_message=error_message,
         )
 
     else:
diff --git a/mitmproxy/tools/console/flowdetailview.py b/mitmproxy/tools/console/flowdetailview.py
index ec716936..fb2494e8 100644
--- a/mitmproxy/tools/console/flowdetailview.py
+++ b/mitmproxy/tools/console/flowdetailview.py
@@ -1,3 +1,4 @@
+import typing
 import urwid
 
 import mitmproxy.flow
@@ -19,6 +20,8 @@ def flowdetails(state, flow: mitmproxy.flow.Flow):
 
     sc = flow.server_conn
     cc = flow.client_conn
+    req: typing.Optional[http.HTTPRequest]
+    resp: typing.Optional[http.HTTPResponse]
     if isinstance(flow, http.HTTPFlow):
         req = flow.request
         resp = flow.response
diff --git a/test/mitmproxy/addons/test_clientplayback.py b/test/mitmproxy/addons/test_clientplayback.py
index 1929ee3d..ecab29df 100644
--- a/test/mitmproxy/addons/test_clientplayback.py
+++ b/test/mitmproxy/addons/test_clientplayback.py
@@ -144,6 +144,9 @@ class TestClientPlayback:
             f.request.raw_content = None
             assert "missing content" in cp.check(f)
 
+            f = tflow.ttcpflow()
+            assert "Can only replay HTTP" in cp.check(f)
+
     @pytest.mark.asyncio
     async def test_playback(self):
         cp = clientplayback.ClientPlayback()
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
index f5088a68..506924e4 100644
--- a/test/mitmproxy/addons/test_view.py
+++ b/test/mitmproxy/addons/test_view.py
@@ -36,7 +36,7 @@ def test_order_refresh():
         assert sargs
 
 
-def test_order_generators():
+def test_order_generators_http():
     v = view.View()
     tf = tflow.tflow(resp=True)
 
@@ -53,6 +53,23 @@ def test_order_generators():
     assert sz.generate(tf) == len(tf.request.raw_content) + len(tf.response.raw_content)
 
 
+def test_order_generators_tcp():
+    v = view.View()
+    tf = tflow.ttcpflow()
+
+    rs = view.OrderRequestStart(v)
+    assert rs.generate(tf) == 946681200
+
+    rm = view.OrderRequestMethod(v)
+    assert rm.generate(tf) == "TCP"
+
+    ru = view.OrderRequestURL(v)
+    assert ru.generate(tf) == "address:22"
+
+    sz = view.OrderKeySize(v)
+    assert sz.generate(tf) == sum(len(m.content) for m in tf.messages)
+
+
 def test_simple():
     v = view.View()
     f = tft(start=1)
@@ -105,6 +122,21 @@ def test_simple():
     assert len(v._store) == 0
 
 
+def test_simple_tcp():
+    v = view.View()
+    f = tflow.ttcpflow()
+    assert v.store_count() == 0
+    v.tcp_start(f)
+    assert list(v) == [f]
+
+    # These all just call update
+    v.tcp_start(f)
+    v.tcp_message(f)
+    v.tcp_error(f)
+    v.tcp_end(f)
+    assert list(v) == [f]
+
+
 def test_filter():
     v = view.View()
     v.request(tft(method="get"))
diff --git a/test/mitmproxy/test_http.py b/test/mitmproxy/test_http.py
index 8a299d8e..6526b56a 100644
--- a/test/mitmproxy/test_http.py
+++ b/test/mitmproxy/test_http.py
@@ -254,6 +254,10 @@ class TestHTTPFlow:
         f.response.decode()
         assert f.response.raw_content == b"abarb"
 
+    def test_timestamp_start(self):
+        f = tflow.tflow()
+        assert f.timestamp_start == f.request.timestamp_start
+
 
 def test_make_error_response():
     resp = http.make_error_response(543, 'foobar', Headers())
diff --git a/test/mitmproxy/tools/console/test_common.py b/test/mitmproxy/tools/console/test_common.py
index 72438c49..1f59ac4e 100644
--- a/test/mitmproxy/tools/console/test_common.py
+++ b/test/mitmproxy/tools/console/test_common.py
@@ -5,10 +5,16 @@ from mitmproxy.tools.console import common
 
 
 def test_format_flow():
-    f = tflow.tflow(resp=True)
-    assert common.format_flow(f, True)
-    assert common.format_flow(f, True, hostheader=True)
-    assert common.format_flow(f, True, extended=True)
+    flows = [
+        tflow.tflow(resp=True),
+        tflow.tflow(err=True),
+        tflow.ttcpflow(),
+        tflow.ttcpflow(err=True),
+    ]
+    for f in flows:
+        for render_mode in common.RenderMode:
+            assert common.format_flow(f, render_mode=render_mode)
+            assert common.format_flow(f, render_mode=render_mode, hostheader=True, focused=False)
 
 
 def test_format_keyvals():
@@ -26,7 +32,7 @@ def test_format_keyvals():
             )
         ), 1
     )
-    assert wrapped.render((30, ))
+    assert wrapped.render((30,))
     assert common.format_keyvals(
         [
             ("aa", wrapped)