From 8cd5e2d25b6db818518c195af9fcaa62bc46a63a Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Sun, 12 Apr 2020 01:22:44 +0200
Subject: lint, mypy, tests++

---
 test/mitmproxy/addons/test_clientplayback.py |  3 +++
 test/mitmproxy/addons/test_view.py           | 34 +++++++++++++++++++++++++++-
 test/mitmproxy/test_http.py                  |  4 ++++
 test/mitmproxy/tools/console/test_common.py  | 16 +++++++++----
 4 files changed, 51 insertions(+), 6 deletions(-)

(limited to 'test')

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)
-- 
cgit v1.2.3