eventsequence: coverage++

8 files changed, 72 insertions, 96 deletions

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/events.py b/mitmproxy/eventsequence.py
index 53f236ca..905cb7bc 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 ValueError
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/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/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/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/test_eventsequence.py b/test/mitmproxy/test_eventsequence.py
index 262df4b0..6e254225 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(ValueError):
+        next(eventsequence.iterate(42))
diff --git a/tox.ini b/tox.ini
index 2acb1e73..fff4d913 100644
--- a/tox.ini
+++ b/tox.ini
@@ -17,7 +17,7 @@ commands =
     --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/events.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 \
+    --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 ""}