Merge pull request #2003 from Kriechi/coverage++

test refactoring and coverage++

25 files changed, 295 insertions, 224 deletions

diff --git a/mitmproxy/addons/dumper.py b/mitmproxy/addons/dumper.py
index 12b0c34b..222f1167 100644
--- a/mitmproxy/addons/dumper.py
+++ b/mitmproxy/addons/dumper.py
@@ -238,7 +238,7 @@ class Dumper:
     def websocket_message(self, f):
         if self.match(f):
             message = f.messages[-1]
-            self.echo(message.info)
+            self.echo(f.message_info(message))
             if self.flow_detail >= 3:
                 self._echo_message(message)
 
diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py
index 07a8975a..b3a93571 100644
--- a/mitmproxy/addons/script.py
+++ b/mitmproxy/addons/script.py
@@ -8,7 +8,7 @@ import types
 
 from mitmproxy import exceptions
 from mitmproxy import ctx
-from mitmproxy import events
+from mitmproxy import eventsequence
 
 
 import watchdog.events
@@ -141,7 +141,7 @@ class Script:
         self.last_options = None
         self.should_reload = threading.Event()
 
-        for i in events.Events:
+        for i in eventsequence.Events:
             if not hasattr(self, i):
                 def mkprox():
                     evt = i
@@ -211,7 +211,7 @@ class ScriptLoader:
             raise ValueError(str(e))
         sc.load_script()
         for f in flows:
-            for evt, o in events.event_sequence(f):
+            for evt, o in eventsequence.iterate(f):
                 sc.run(evt, o)
         sc.done()
         return sc
diff --git a/mitmproxy/connections.py b/mitmproxy/connections.py
index 9c4bca2f..a32889bd 100644
--- a/mitmproxy/connections.py
+++ b/mitmproxy/connections.py
@@ -1,6 +1,5 @@
 import time
 
-import copy
 import os
 
 from mitmproxy import stateobject
@@ -82,9 +81,6 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
         tls_version=str,
     )
 
-    def copy(self):
-        return copy.copy(self)
-
     def send(self, message):
         if isinstance(message, list):
             message = b''.join(message)
@@ -222,9 +218,6 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
             via=None
         ))
 
-    def copy(self):
-        return copy.copy(self)
-
     def connect(self):
         self.timestamp_start = time.time()
         tcp.TCPClient.connect(self)
diff --git a/mitmproxy/events.py b/mitmproxy/eventsequence.py
index 53f236ca..5872f607 100644
--- a/mitmproxy/events.py
+++ b/mitmproxy/eventsequence.py
@@ -37,7 +37,7 @@ Events = frozenset([
 ])
 
 
-def event_sequence(f):
+def iterate(f):
     if isinstance(f, http.HTTPFlow):
         if f.request:
             yield "requestheaders", f
@@ -70,4 +70,4 @@ def event_sequence(f):
             yield "tcp_error", f
         yield "tcp_end", f
     else:
-        raise NotImplementedError
+        raise TypeError()
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 7034cb4a..5ef957c9 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -1,5 +1,4 @@
 import time
-import copy
 import uuid
 
 from mitmproxy import controller  # noqa
@@ -7,7 +6,7 @@ from mitmproxy import stateobject
 from mitmproxy import connections
 from mitmproxy import version
 
-import typing # noqa
+import typing  # noqa
 
 
 class Error(stateobject.StateObject):
@@ -53,10 +52,6 @@ class Error(stateobject.StateObject):
         f.set_state(state)
         return f
 
-    def copy(self):
-        c = copy.copy(self)
-        return c
-
 
 class Flow(stateobject.StateObject):
 
@@ -116,16 +111,9 @@ class Flow(stateobject.StateObject):
         return f
 
     def copy(self):
-        f = copy.copy(self)
-
+        f = super().copy()
         f.id = str(uuid.uuid4())
         f.live = False
-        f.client_conn = self.client_conn.copy()
-        f.server_conn = self.server_conn.copy()
-        f.metadata = self.metadata.copy()
-
-        if self.error:
-            f.error = self.error.copy()
         return f
 
     def modified(self):
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index ee240eeb..3a3f4399 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -7,7 +7,7 @@ import sys
 from mitmproxy import addonmanager
 from mitmproxy import options
 from mitmproxy import controller
-from mitmproxy import events
+from mitmproxy import eventsequence
 from mitmproxy import exceptions
 from mitmproxy import connections
 from mitmproxy import http
@@ -91,7 +91,7 @@ class Master:
         changed = False
         try:
             mtype, obj = self.event_queue.get(timeout=timeout)
-            if mtype not in events.Events:
+            if mtype not in eventsequence.Events:
                 raise exceptions.ControlException(
                     "Unknown event %s" % repr(mtype)
                 )
@@ -153,7 +153,7 @@ class Master:
                 f.request.port = self.server.config.upstream_server.address.port
                 f.request.scheme = self.server.config.upstream_server.scheme
         f.reply = controller.DummyReply()
-        for e, o in events.event_sequence(f):
+        for e, o in eventsequence.iterate(f):
             getattr(self, e)(o)
 
     def load_flows(self, fr: io.FlowReader) -> int:
diff --git a/mitmproxy/net/http/http1/read.py b/mitmproxy/net/http/http1/read.py
index 6eb30709..d0493da4 100644
--- a/mitmproxy/net/http/http1/read.py
+++ b/mitmproxy/net/http/http1/read.py
@@ -158,8 +158,9 @@ def connection_close(http_version, headers):
     """
         Checks the message to see if the client connection should be closed
         according to RFC 2616 Section 8.1.
+        If we don't have a Connection header, HTTP 1.1 connections are assumed
+        to be persistent.
     """
-    # At first, check if we have an explicit Connection header.
     if "connection" in headers:
         tokens = get_header_tokens(headers, "connection")
         if "close" in tokens:
@@ -167,9 +168,7 @@ def connection_close(http_version, headers):
         elif "keep-alive" in tokens:
             return False
 
-    # If we don't have a Connection header, HTTP 1.1 connections are assumed to
-    # be persistent
-    return http_version != "HTTP/1.1" and http_version != b"HTTP/1.1"  # FIXME: Remove one case.
+    return http_version != "HTTP/1.1" and http_version != b"HTTP/1.1"
 
 
 def expected_http_body_size(request, response=None):
diff --git a/mitmproxy/net/http/http2/utils.py b/mitmproxy/net/http/http2/utils.py
index 24dc773c..4a553d8d 100644
--- a/mitmproxy/net/http/http2/utils.py
+++ b/mitmproxy/net/http/http2/utils.py
@@ -21,7 +21,6 @@ def parse_headers(headers):
         first_line_format = "relative"
     else:
         first_line_format = "absolute"
-        # FIXME: verify if path or :host contains what we need
         scheme, host, port, _ = url.parse(path)
 
     if authority:
diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py
index d0b12540..e170f19d 100644
--- a/mitmproxy/proxy/protocol/websocket.py
+++ b/mitmproxy/proxy/protocol/websocket.py
@@ -8,7 +8,7 @@ from mitmproxy import flow
 from mitmproxy.proxy.protocol import base
 from mitmproxy.net import tcp
 from mitmproxy.net import websockets
-from mitmproxy.websocket import WebSocketFlow, WebSocketBinaryMessage, WebSocketTextMessage
+from mitmproxy.websocket import WebSocketFlow, WebSocketMessage
 
 
 class WebSocketLayer(base.Layer):
@@ -65,12 +65,7 @@ class WebSocketLayer(base.Layer):
             compressed_message = fb[0].header.rsv1
             fb.clear()
 
-            if message_type == websockets.OPCODE.TEXT:
-                t = WebSocketTextMessage
-            else:
-                t = WebSocketBinaryMessage
-
-            websocket_message = t(self.flow, not is_server, payload)
+            websocket_message = WebSocketMessage(message_type, not is_server, payload)
             length = len(websocket_message.content)
             self.flow.messages.append(websocket_message)
             self.channel.ask("websocket_message", self.flow)
diff --git a/mitmproxy/script/concurrent.py b/mitmproxy/script/concurrent.py
index 2fd7ad8d..366929a5 100644
--- a/mitmproxy/script/concurrent.py
+++ b/mitmproxy/script/concurrent.py
@@ -3,7 +3,7 @@ This module provides a @concurrent decorator primitive to
 offload computations from mitmproxy's main master thread.
 """
 
-from mitmproxy import events
+from mitmproxy import eventsequence
 from mitmproxy.types import basethread
 
 
@@ -12,7 +12,7 @@ class ScriptThread(basethread.BaseThread):
 
 
 def concurrent(fn):
-    if fn.__name__ not in events.Events - {"start", "configure", "tick"}:
+    if fn.__name__ not in eventsequence.Events - {"start", "configure", "tick"}:
         raise NotImplementedError(
             "Concurrent decorator not supported for '%s' method." % fn.__name__
         )
diff --git a/mitmproxy/tcp.py b/mitmproxy/tcp.py
index 3f10f82b..067fbfe3 100644
--- a/mitmproxy/tcp.py
+++ b/mitmproxy/tcp.py
@@ -9,8 +9,8 @@ from mitmproxy.types import serializable
 class TCPMessage(serializable.Serializable):
 
     def __init__(self, from_client, content, timestamp=None):
-        self.content = content
         self.from_client = from_client
+        self.content = content
         self.timestamp = timestamp or time.time()
 
     @classmethod
@@ -21,9 +21,7 @@ class TCPMessage(serializable.Serializable):
         return self.from_client, self.content, self.timestamp
 
     def set_state(self, state):
-        self.from_client = state.pop("from_client")
-        self.content = state.pop("content")
-        self.timestamp = state.pop("timestamp")
+        self.from_client, self.content, self.timestamp = state
 
     def __repr__(self):
         return "{direction} {content}".format(
diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py
index a25b6891..bb8daa02 100644
--- a/mitmproxy/test/taddons.py
+++ b/mitmproxy/test/taddons.py
@@ -3,7 +3,7 @@ import contextlib
 import mitmproxy.master
 import mitmproxy.options
 from mitmproxy import proxy
-from mitmproxy import events
+from mitmproxy import eventsequence
 from mitmproxy import exceptions
 
 
@@ -57,7 +57,7 @@ class context:
             is taken (as in flow interception).
         """
         f.reply._state = "handled"
-        for evt, arg in events.event_sequence(f):
+        for evt, arg in eventsequence.iterate(f):
             h = getattr(addon, evt, None)
             if h:
                 h(arg)
diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py
index a5670538..6d330840 100644
--- a/mitmproxy/test/tflow.py
+++ b/mitmproxy/test/tflow.py
@@ -1,3 +1,4 @@
+from mitmproxy.net import websockets
 from mitmproxy.test import tutils
 from mitmproxy import tcp
 from mitmproxy import websocket
@@ -70,8 +71,8 @@ def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None,
 
     if messages is True:
         messages = [
-            websocket.WebSocketBinaryMessage(f, True, b"hello binary"),
-            websocket.WebSocketTextMessage(f, False, "hello text".encode()),
+            websocket.WebSocketMessage(websockets.OPCODE.BINARY, True, b"hello binary"),
+            websocket.WebSocketMessage(websockets.OPCODE.TEXT, False, "hello text".encode()),
         ]
     if err is True:
         err = terr()
diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py
index 6e998a52..25a82878 100644
--- a/mitmproxy/websocket.py
+++ b/mitmproxy/websocket.py
@@ -1,63 +1,38 @@
 import time
-
-from typing import List
+from typing import List, Optional
 
 from mitmproxy import flow
 from mitmproxy.http import HTTPFlow
 from mitmproxy.net import websockets
-from mitmproxy.utils import strutils
 from mitmproxy.types import serializable
+from mitmproxy.utils import strutils
 
 
 class WebSocketMessage(serializable.Serializable):
-
-    def __init__(self, flow, from_client, content, timestamp=None):
-        self.flow = flow
-        self.content = content
+    def __init__(self, type: int, from_client: bool, content: bytes, timestamp: Optional[int]=None):
+        self.type = type
         self.from_client = from_client
-        self.timestamp = timestamp or time.time()
+        self.content = content
+        self.timestamp = timestamp or time.time()  # type: int
 
     @classmethod
     def from_state(cls, state):
         return cls(*state)
 
     def get_state(self):
-        return self.from_client, self.content, self.timestamp
+        return self.type, self.from_client, self.content, self.timestamp
 
     def set_state(self, state):
-        self.from_client = state.pop("from_client")
-        self.content = state.pop("content")
-        self.timestamp = state.pop("timestamp")
-
-    @property
-    def info(self):
-        return "{client} {direction} WebSocket {type} message {direction} {server}{endpoint}".format(
-            type=self.type,
-            client=repr(self.flow.client_conn.address),
-            server=repr(self.flow.server_conn.address),
-            direction="->" if self.from_client else "<-",
-            endpoint=self.flow.handshake_flow.request.path,
-        )
-
-
-class WebSocketBinaryMessage(WebSocketMessage):
-
-    type = 'binary'
+        self.type, self.from_client, self.content, self.timestamp = state
 
     def __repr__(self):
-        return "binary message: {}".format(strutils.bytes_to_escaped_str(self.content))
-
-
-class WebSocketTextMessage(WebSocketMessage):
-
-    type = 'text'
-
-    def __repr__(self):
-        return "text message: {}".format(repr(self.content))
+        if self.type == websockets.OPCODE.TEXT:
+            return "text message: {}".format(repr(self.content))
+        else:
+            return "binary message: {}".format(strutils.bytes_to_escaped_str(self.content))
 
 
 class WebSocketFlow(flow.Flow):
-
     """
     A WebsocketFlow is a simplified representation of a Websocket session.
     """
@@ -70,18 +45,55 @@ class WebSocketFlow(flow.Flow):
         self.close_message = '(message missing)'
         self.close_reason = 'unknown status code'
         self.handshake_flow = handshake_flow
-        self.client_key = websockets.get_client_key(self.handshake_flow.request.headers)
-        self.client_protocol = websockets.get_protocol(self.handshake_flow.request.headers)
-        self.client_extensions = websockets.get_extensions(self.handshake_flow.request.headers)
-        self.server_accept = websockets.get_server_accept(self.handshake_flow.response.headers)
-        self.server_protocol = websockets.get_protocol(self.handshake_flow.response.headers)
-        self.server_extensions = websockets.get_extensions(self.handshake_flow.response.headers)
 
     _stateobject_attributes = flow.Flow._stateobject_attributes.copy()
     _stateobject_attributes.update(
         messages=List[WebSocketMessage],
+        close_sender=str,
+        close_code=str,
+        close_message=str,
+        close_reason=str,
         handshake_flow=HTTPFlow,
     )
 
+    @classmethod
+    def from_state(cls, state):
+        f = cls(None, None, None)
+        f.set_state(state)
+        return f
+
     def __repr__(self):
-        return "WebSocketFlow ({} messages)".format(len(self.messages))
+        return "<WebSocketFlow ({} messages)>".format(len(self.messages))
+
+    @property
+    def client_key(self):
+        return websockets.get_client_key(self.handshake_flow.request.headers)
+
+    @property
+    def client_protocol(self):
+        return websockets.get_protocol(self.handshake_flow.request.headers)
+
+    @property
+    def client_extensions(self):
+        return websockets.get_extensions(self.handshake_flow.request.headers)
+
+    @property
+    def server_accept(self):
+        return websockets.get_server_accept(self.handshake_flow.response.headers)
+
+    @property
+    def server_protocol(self):
+        return websockets.get_protocol(self.handshake_flow.response.headers)
+
+    @property
+    def server_extensions(self):
+        return websockets.get_extensions(self.handshake_flow.response.headers)
+
+    def message_info(self, message: WebSocketMessage) -> str:
+        return "{client} {direction} WebSocket {type} message {direction} {server}{endpoint}".format(
+            type=message.type,
+            client=repr(self.client_conn.address),
+            server=repr(self.server_conn.address),
+            direction="->" if message.from_client else "<-",
+            endpoint=self.handshake_flow.request.path,
+        )
diff --git a/test/conftest.py b/test/conftest.py
index 50ec3421..83823a19 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -103,15 +103,17 @@ def pytest_runtestloop(session):
     measured_files = [os.path.normpath(os.path.relpath(f, prefix)) for f in cov.get_data().measured_files()]
     measured_files = [f for f in measured_files if not any(f.startswith(excluded_f) for excluded_f in excluded_files)]
 
-    for name in pytest.config.option.full_cov:
+    for name in coverage_values.keys():
         files = [f for f in measured_files if f.startswith(os.path.normpath(name))]
         try:
             with open(os.devnull, 'w') as null:
-                coverage_values[name] = cov.report(files, ignore_errors=True, file=null)
+                overall = cov.report(files, ignore_errors=True, file=null)
+                singles = [(s, cov.report(s, ignore_errors=True, file=null)) for s in files]
+                coverage_values[name] = (overall, singles)
         except:
             pass
 
-    if any(v < 100 for v in coverage_values.values()):
+    if any(v < 100 for v, _ in coverage_values.values()):
         # make sure we get the EXIT_TESTSFAILED exit code
         session.testsfailed += 1
     else:
@@ -132,12 +134,15 @@ def pytest_terminal_summary(terminalreporter, exitstatus):
         msg = "FAIL: Full test coverage not reached!\n"
         terminalreporter.write(msg, **markup)
 
-        for name, value in coverage_values.items():
-            if value < 100:
+        for name in sorted(coverage_values.keys()):
+            msg = 'Coverage for {}: {:.2f}%\n'.format(name, coverage_values[name][0])
+            if coverage_values[name][0] < 100:
                 markup = {'red': True, 'bold': True}
+                for s, v in sorted(coverage_values[name][1]):
+                    if v < 100:
+                        msg += '  {}: {:.2f}%\n'.format(s, v)
             else:
                 markup = {'green': True}
-            msg = 'Coverage for {}: {:.2f}%\n'.format(name, value)
             terminalreporter.write(msg, **markup)
     else:
         markup = {'green': True}
@@ -145,6 +150,7 @@ def pytest_terminal_summary(terminalreporter, exitstatus):
         msg += '{}\n\n'.format('\n'.join(pytest.config.option.full_cov))
         terminalreporter.write(msg, **markup)
 
-    msg = 'Excluded files:\n'
-    msg += '{}\n'.format('\n'.join(pytest.config.option.no_full_cov))
+    msg = '\nExcluded files:\n'
+    for s in sorted(pytest.config.option.no_full_cov):
+        msg += "  {}\n".format(s)
     terminalreporter.write(msg)
diff --git a/test/mitmproxy/addons/test_stickycookie.py b/test/mitmproxy/addons/test_stickycookie.py
index 157f2959..9092e09b 100644
--- a/test/mitmproxy/addons/test_stickycookie.py
+++ b/test/mitmproxy/addons/test_stickycookie.py
@@ -39,7 +39,6 @@ class TestStickyCookie:
             assert "cookie" not in f.request.headers
 
             f = f.copy()
-            f.reply.acked = False
             sc.request(f)
             assert f.request.headers["cookie"] == "foo=bar"
 
diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py
index 5be88e5c..6b9b6ea8 100644
--- a/test/mitmproxy/data/addonscripts/recorder.py
+++ b/test/mitmproxy/data/addonscripts/recorder.py
@@ -1,5 +1,5 @@
 from mitmproxy import controller
-from mitmproxy import events
+from mitmproxy import eventsequence
 from mitmproxy import ctx
 import sys
 
@@ -11,7 +11,7 @@ class CallLogger:
         self.name = name
 
     def __getattr__(self, attr):
-        if attr in events.Events:
+        if attr in eventsequence.Events:
             def prox(*args, **kwargs):
                 lg = (self.name, attr, args, kwargs)
                 if attr != "log":
diff --git a/test/mitmproxy/protocol/test_websocket.py b/test/mitmproxy/protocol/test_websocket.py
index e42250e0..73ee8b35 100644
--- a/test/mitmproxy/protocol/test_websocket.py
+++ b/test/mitmproxy/protocol/test_websocket.py
@@ -179,16 +179,15 @@ class TestSimple(_WebSocketTest):
         assert isinstance(self.master.state.flows[1], WebSocketFlow)
         assert len(self.master.state.flows[1].messages) == 5
         assert self.master.state.flows[1].messages[0].content == b'server-foobar'
-        assert self.master.state.flows[1].messages[0].type == 'text'
+        assert self.master.state.flows[1].messages[0].type == websockets.OPCODE.TEXT
         assert self.master.state.flows[1].messages[1].content == b'client-foobar'
-        assert self.master.state.flows[1].messages[1].type == 'text'
+        assert self.master.state.flows[1].messages[1].type == websockets.OPCODE.TEXT
         assert self.master.state.flows[1].messages[2].content == b'client-foobar'
-        assert self.master.state.flows[1].messages[2].type == 'text'
+        assert self.master.state.flows[1].messages[2].type == websockets.OPCODE.TEXT
         assert self.master.state.flows[1].messages[3].content == b'\xde\xad\xbe\xef'
-        assert self.master.state.flows[1].messages[3].type == 'binary'
+        assert self.master.state.flows[1].messages[3].type == websockets.OPCODE.BINARY
         assert self.master.state.flows[1].messages[4].content == b'\xde\xad\xbe\xef'
-        assert self.master.state.flows[1].messages[4].type == 'binary'
-        assert [m.info for m in self.master.state.flows[1].messages]
+        assert self.master.state.flows[1].messages[4].type == websockets.OPCODE.BINARY
 
 
 class TestSimpleTLS(_WebSocketTest):
diff --git a/test/mitmproxy/test_eventsequence.py b/test/mitmproxy/test_eventsequence.py
index 262df4b0..fe0f92b3 100644
--- a/test/mitmproxy/test_eventsequence.py
+++ b/test/mitmproxy/test_eventsequence.py
@@ -1,81 +1,57 @@
-from mitmproxy import events
-import contextlib
-from . import tservers
-
-
-class Eventer:
-    def __init__(self, **handlers):
-        self.failure = None
-        self.called = []
-        self.handlers = handlers
-        for i in events.Events - {"tick"}:
-            def mkprox():
-                evt = i
-
-                def prox(*args, **kwargs):
-                    self.called.append(evt)
-                    if evt in self.handlers:
-                        try:
-                            handlers[evt](*args, **kwargs)
-                        except AssertionError as e:
-                            self.failure = e
-                return prox
-            setattr(self, i, mkprox())
-
-    def fail(self):
-        pass
-
-
-class SequenceTester:
-    @contextlib.contextmanager
-    def addon(self, addon):
-        self.master.addons.add(addon)
-        yield
-        self.master.addons.remove(addon)
-        if addon.failure:
-            raise addon.failure
-
-
-class TestBasic(tservers.HTTPProxyTest, SequenceTester):
-    ssl = True
-
-    def test_requestheaders(self):
-
-        def hdrs(f):
-            assert f.request.headers
-            assert not f.request.content
-
-        def req(f):
-            assert f.request.headers
-            assert f.request.content
-
-        with self.addon(Eventer(requestheaders=hdrs, request=req)):
-            p = self.pathoc()
-            with p.connect():
-                assert p.request("get:'/p/200':b@10").status_code == 200
-
-    def test_100_continue_fail(self):
-        e = Eventer()
-        with self.addon(e):
-            p = self.pathoc()
-            with p.connect():
-                p.request(
-                    """
-                        get:'/p/200'
-                        h'expect'='100-continue'
-                        h'content-length'='1000'
-                        da
-                    """
-                )
-        assert "requestheaders" in e.called
-        assert "responseheaders" not in e.called
-
-    def test_connect(self):
-        e = Eventer()
-        with self.addon(e):
-            p = self.pathoc()
-            with p.connect():
-                p.request("get:'/p/200:b@1'")
-            assert "http_connect" in e.called
-            assert e.called.count("requestheaders") == 1
-            assert e.called.count("request") == 1
+import pytest
+
+from mitmproxy import eventsequence
+from mitmproxy.test import tflow
+
+
+@pytest.mark.parametrize("resp, err", [
+    (False, False),
+    (True, False),
+    (False, True),
+    (True, True),
+])
+def test_http_flow(resp, err):
+    f = tflow.tflow(resp=resp, err=err)
+    i = eventsequence.iterate(f)
+    assert next(i) == ("requestheaders", f)
+    assert next(i) == ("request", f)
+    if resp:
+        assert next(i) == ("responseheaders", f)
+        assert next(i) == ("response", f)
+    if err:
+        assert next(i) == ("error", f)
+
+
+@pytest.mark.parametrize("err", [False, True])
+def test_websocket_flow(err):
+    f = tflow.twebsocketflow(err=err)
+    i = eventsequence.iterate(f)
+    assert next(i) == ("websocket_start", f)
+    assert len(f.messages) == 0
+    assert next(i) == ("websocket_message", f)
+    assert len(f.messages) == 1
+    assert next(i) == ("websocket_message", f)
+    assert len(f.messages) == 2
+    if err:
+        assert next(i) == ("websocket_error", f)
+    assert next(i) == ("websocket_end", f)
+
+
+@pytest.mark.parametrize("err", [False, True])
+def test_tcp_flow(err):
+    f = tflow.ttcpflow(err=err)
+    i = eventsequence.iterate(f)
+    assert next(i) == ("tcp_start", f)
+    assert len(f.messages) == 0
+    assert next(i) == ("tcp_message", f)
+    assert len(f.messages) == 1
+    assert next(i) == ("tcp_message", f)
+    assert len(f.messages) == 2
+    if err:
+        assert next(i) == ("tcp_error", f)
+    assert next(i) == ("tcp_end", f)
+
+
+def test_invalid():
+    with pytest.raises(TypeError):
+        next(eventsequence.iterate(42))
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 3b5d0ac1..65e6845f 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -10,6 +10,7 @@ from mitmproxy.exceptions import FlowReadException, Kill
 from mitmproxy import flow
 from mitmproxy import http
 from mitmproxy import connections
+from mitmproxy import tcp
 from mitmproxy.proxy import ProxyConfig
 from mitmproxy.proxy.server import DummyServer
 from mitmproxy import master
@@ -156,8 +157,99 @@ class TestHTTPFlow:
         assert f.response.raw_content == b"abarb"
 
 
+class TestWebSocketFlow:
+
+    def test_copy(self):
+        f = tflow.twebsocketflow()
+        f.get_state()
+        f2 = f.copy()
+        a = f.get_state()
+        b = f2.get_state()
+        del a["id"]
+        del b["id"]
+        del a["handshake_flow"]["id"]
+        del b["handshake_flow"]["id"]
+        assert a == b
+        assert not f == f2
+        assert f is not f2
+
+        assert f.client_key == f2.client_key
+        assert f.client_protocol == f2.client_protocol
+        assert f.client_extensions == f2.client_extensions
+        assert f.server_accept == f2.server_accept
+        assert f.server_protocol == f2.server_protocol
+        assert f.server_extensions == f2.server_extensions
+        assert f.messages is not f2.messages
+        assert f.handshake_flow is not f2.handshake_flow
+
+        for m in f.messages:
+            m2 = m.copy()
+            m2.set_state(m2.get_state())
+            assert m is not m2
+            assert m.get_state() == m2.get_state()
+
+        f = tflow.twebsocketflow(err=True)
+        f2 = f.copy()
+        assert f is not f2
+        assert f.handshake_flow is not f2.handshake_flow
+        assert f.error.get_state() == f2.error.get_state()
+        assert f.error is not f2.error
+
+    def test_match(self):
+        f = tflow.twebsocketflow()
+        assert not flowfilter.match("~b nonexistent", f)
+        assert flowfilter.match(None, f)
+        assert not flowfilter.match("~b nonexistent", f)
+
+        f = tflow.twebsocketflow(err=True)
+        assert flowfilter.match("~e", f)
+
+        with pytest.raises(ValueError):
+            flowfilter.match("~", f)
+
+    def test_repr(self):
+        f = tflow.twebsocketflow()
+        assert 'WebSocketFlow' in repr(f)
+        assert 'binary message: ' in repr(f.messages[0])
+        assert 'text message: ' in repr(f.messages[1])
+
+
 class TestTCPFlow:
 
+    def test_copy(self):
+        f = tflow.ttcpflow()
+        f.get_state()
+        f2 = f.copy()
+        a = f.get_state()
+        b = f2.get_state()
+        del a["id"]
+        del b["id"]
+        assert a == b
+        assert not f == f2
+        assert f is not f2
+
+        assert f.messages is not f2.messages
+
+        for m in f.messages:
+            assert m.get_state()
+            m2 = m.copy()
+            assert not m == m2
+            assert m is not m2
+
+            a = m.get_state()
+            b = m2.get_state()
+            assert a == b
+
+        m = tcp.TCPMessage(False, 'foo')
+        m.set_state(f.messages[0].get_state())
+        assert m.timestamp == f.messages[0].timestamp
+
+        f = tflow.ttcpflow(err=True)
+        f2 = f.copy()
+        assert f is not f2
+        assert f.error.get_state() == f2.error.get_state()
+        assert f.error is not f2.error
+
     def test_match(self):
         f = tflow.ttcpflow()
         assert not flowfilter.match("~b nonexistent", f)
@@ -170,6 +262,11 @@ class TestTCPFlow:
         with pytest.raises(ValueError):
             flowfilter.match("~", f)
 
+    def test_repr(self):
+        f = tflow.ttcpflow()
+        assert 'TCPFlow' in repr(f)
+        assert '-> ' in repr(f.messages[0])
+
 
 class TestSerialize:
 
@@ -308,11 +405,11 @@ class TestFlowMaster:
         fm.clientconnect(f.client_conn)
         f.request = http.HTTPRequest.wrap(mitmproxy.test.tutils.treq())
         fm.request(f)
-        assert s.flow_count() == 1
+        assert len(s.flows) == 1
 
         f.response = http.HTTPResponse.wrap(mitmproxy.test.tutils.tresp())
         fm.response(f)
-        assert s.flow_count() == 1
+        assert len(s.flows) == 1
 
         fm.clientdisconnect(f.client_conn)
 
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index dfd61540..9cd47cea 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -984,16 +984,16 @@ class TestUpstreamProxySSL(
         assert req.status_code == 418
 
         # CONNECT from pathoc to chain[0],
-        assert self.proxy.tmaster.state.flow_count() == 1
+        assert len(self.proxy.tmaster.state.flows) == 1
         assert self.proxy.tmaster.state.flows[0].server_conn.via
         # request from pathoc to chain[0]
         # CONNECT from proxy to chain[1],
-        assert self.chain[0].tmaster.state.flow_count() == 1
+        assert len(self.chain[0].tmaster.state.flows) == 1
         assert self.chain[0].tmaster.state.flows[0].server_conn.via
         # request from proxy to chain[1]
         # request from chain[0] (regular proxy doesn't store CONNECTs)
         assert not self.chain[1].tmaster.state.flows[0].server_conn.via
-        assert self.chain[1].tmaster.state.flow_count() == 1
+        assert len(self.chain[1].tmaster.state.flows) == 1
 
     def test_change_upstream_proxy_connect(self):
         # skip chain[0].
@@ -1050,17 +1050,17 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
             assert req.status_code == 418
 
             # First request goes through all three proxies exactly once
-            assert self.proxy.tmaster.state.flow_count() == 1
-            assert self.chain[0].tmaster.state.flow_count() == 1
-            assert self.chain[1].tmaster.state.flow_count() == 1
+            assert len(self.proxy.tmaster.state.flows) == 1
+            assert len(self.chain[0].tmaster.state.flows) == 1
+            assert len(self.chain[1].tmaster.state.flows) == 1
 
             req = p.request("get:'/p/418:b\"content2\"'")
             assert req.status_code == 502
 
-            assert self.proxy.tmaster.state.flow_count() == 2
-            assert self.chain[0].tmaster.state.flow_count() == 2
+            assert len(self.proxy.tmaster.state.flows) == 2
+            assert len(self.chain[0].tmaster.state.flows) == 2
             # Upstream sees two requests due to reconnection attempt
-            assert self.chain[1].tmaster.state.flow_count() == 3
+            assert len(self.chain[1].tmaster.state.flows) == 3
             assert not self.chain[1].tmaster.state.flows[-1].response
             assert not self.chain[1].tmaster.state.flows[-2].response
 
diff --git a/test/mitmproxy/test_stateobject.py b/test/mitmproxy/test_stateobject.py
index edec92c2..7b8e30d0 100644
--- a/test/mitmproxy/test_stateobject.py
+++ b/test/mitmproxy/test_stateobject.py
@@ -1,4 +1,5 @@
 from typing import List
+import pytest
 
 from mitmproxy.stateobject import StateObject
 
@@ -67,3 +68,14 @@ def test_container_list():
     assert len(copy.children) == 2
     assert copy.children is not a.children
     assert copy.children[0] is not a.children[0]
+
+
+def test_too_much_state():
+    a = Container()
+    a.child = Child(42)
+    s = a.get_state()
+    s['foo'] = 'bar'
+    b = Container()
+
+    with pytest.raises(RuntimeWarning):
+        b.set_state(s)
diff --git a/test/mitmproxy/test_version.py b/test/mitmproxy/test_version.py
new file mode 100644
index 00000000..f87b0851
--- /dev/null
+++ b/test/mitmproxy/test_version.py
@@ -0,0 +1,10 @@
+import runpy
+
+from mitmproxy import version
+
+
+def test_version(capsys):
+    runpy.run_module('mitmproxy.version', run_name='__main__')
+    stdout, stderr = capsys.readouterr()
+    assert len(stdout) > 0
+    assert stdout.strip() == version.VERSION
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 170a4917..c9a7e595 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -35,10 +35,6 @@ class TestState:
     #     if f not in self.flows:
     #         self.flows.append(f)
 
-    # FIXME: compat with old state - remove in favor of len(state.flows)
-    def flow_count(self):
-        return len(self.flows)
-
 
 class TestMaster(master.Master):
 
diff --git a/tox.ini b/tox.ini
index a91f5d1a..fff4d913 100644
--- a/tox.ini
+++ b/tox.ini
@@ -12,21 +12,12 @@ setenv = HOME = {envtmpdir}
 commands =
   mitmdump --version
   pytest --timeout 60 --cov-report='' --cov=mitmproxy --cov=pathod \
-    --full-cov=mitmproxy/addons/ \
-    --full-cov=mitmproxy/contentviews/ --no-full-cov=mitmproxy/contentviews/__init__.py --no-full-cov=mitmproxy/contentviews/protobuf.py --no-full-cov=mitmproxy/contentviews/wbxml.py --no-full-cov=mitmproxy/contentviews/xml_html.py \
-    --full-cov=mitmproxy/net/ --no-full-cov=mitmproxy/net/tcp.py --no-full-cov=mitmproxy/net/http/cookies.py --no-full-cov=mitmproxy/net/http/encoding.py --no-full-cov=mitmproxy/net/http/message.py --no-full-cov=mitmproxy/net/http/request.py --no-full-cov=mitmproxy/net/http/response.py --no-full-cov=mitmproxy/net/http/url.py \
-    --full-cov=mitmproxy/proxy/ --no-full-cov=mitmproxy/proxy/protocol/ --no-full-cov=mitmproxy/proxy/config.py --no-full-cov=mitmproxy/proxy/root_context.py --no-full-cov=mitmproxy/proxy/server.py \
-    --full-cov=mitmproxy/script/ \
-    --full-cov=mitmproxy/test/ \
-    --full-cov=mitmproxy/types/ \
-    --full-cov=mitmproxy/utils/ \
-    --full-cov=mitmproxy/__init__.py \
-    --full-cov=mitmproxy/addonmanager.py \
-    --full-cov=mitmproxy/ctx.py \
-    --full-cov=mitmproxy/exceptions.py \
-    --full-cov=mitmproxy/io.py \
-    --full-cov=mitmproxy/log.py \
-    --full-cov=mitmproxy/options.py \
+    --full-cov=mitmproxy/ \
+    --no-full-cov=mitmproxy/contentviews/__init__.py --no-full-cov=mitmproxy/contentviews/protobuf.py --no-full-cov=mitmproxy/contentviews/wbxml.py --no-full-cov=mitmproxy/contentviews/xml_html.py \
+    --no-full-cov=mitmproxy/net/tcp.py --no-full-cov=mitmproxy/net/http/cookies.py --no-full-cov=mitmproxy/net/http/encoding.py --no-full-cov=mitmproxy/net/http/message.py --no-full-cov=mitmproxy/net/http/request.py --no-full-cov=mitmproxy/net/http/response.py --no-full-cov=mitmproxy/net/http/url.py \
+    --no-full-cov=mitmproxy/proxy/protocol/ --no-full-cov=mitmproxy/proxy/config.py --no-full-cov=mitmproxy/proxy/root_context.py --no-full-cov=mitmproxy/proxy/server.py \
+    --no-full-cov=mitmproxy/tools/ \
+    --no-full-cov=mitmproxy/certs.py --no-full-cov=mitmproxy/connections.py --no-full-cov=mitmproxy/controller.py --no-full-cov=mitmproxy/export.py --no-full-cov=mitmproxy/flow.py --no-full-cov=mitmproxy/flowfilter.py --no-full-cov=mitmproxy/http.py --no-full-cov=mitmproxy/io_compat.py --no-full-cov=mitmproxy/master.py --no-full-cov=mitmproxy/optmanager.py \
     --full-cov=pathod/ --no-full-cov=pathod/pathoc.py --no-full-cov=pathod/pathod.py --no-full-cov=pathod/test.py --no-full-cov=pathod/protocols/http2.py \
     {posargs}
   {env:CI_COMMANDS:python -c ""}