protocols: docs++, minor test improvements

mitmproxy/proxy/protocol2/test/test_tcp.py

@@ -88,36 +88,6 @@ def test_simple_explicit(tctx):
     assert list(layer.handle_event(events.HookReply(tcp_end, None))) == []
 
 
-r'''
-def test_simple_alternate_syntax(tctx):
-    """
-    Some alternate syntax experimentations:
-        - no asserts, evaluate when we reach a command or the end.
-        - use <= for final statement
-        - If the final statement is a hook, its data is returned.
-          This replaces placeholders (we must do partial matching on the first <= hook though)
-    """
-    playbook = tutils.playbook(tcp.TCPLayer(tctx))
-
-    flow = (playbook
-        << commands.Hook("tcp_start", mock.Mock())
-        >> events.HookReply(-1, None)
-        << commands.OpenConnection(tctx.server)
-        >> events.OpenConnectionReply(-1, None)
-        >> events.DataReceived(tctx.client, b"hello!")
-        <= commands.Hook("tcp_message", None))
-    assert flow.messages[0].content == b"hello!"
-    (playbook
-        >> events.HookReply(-1, None)
-        << commands.SendData(tctx.server, b"hello!")
-        >> events.ConnectionClosed(tctx.server)
-        << commands.CloseConnection(tctx.client)
-        << commands.Hook("tcp_end", flow)
-        >> events.HookReply(-1, None)
-        <= None)
-'''
-
-
 def test_receive_data_before_server_connected(tctx):
     """
     assert that data received before a server connection is established
@@ -135,18 +105,26 @@ def test_receive_data_before_server_connected(tctx):
         >> events.HookReply(-1, None)
         << commands.SendData(tctx.server, b"hello!")
     )
+    assert f().messages
 
 
 def test_receive_data_after_server_disconnected(tctx):
     """
-    this should just be discarded.
+    data received after a connection has been closed should just be discarded.
     """
+    f = tutils.Placeholder()
     assert (
-        tutils.playbook(tcp.TCPLayer(tctx, True))
+        tutils.playbook(tcp.TCPLayer(tctx))
+        << commands.Hook("tcp_start", f)
+        >> events.HookReply(-1, None)
         << commands.OpenConnection(tctx.server)
         >> events.OpenConnectionReply(-1, None)
         >> events.ConnectionClosed(tctx.server)
         << commands.CloseConnection(tctx.client)
+        << commands.Hook("tcp_end", f)
+        >> events.HookReply(-1, None)
         >> events.DataReceived(tctx.client, b"i'm late")
         << None
     )
+    # not included here as it has not been sent to the server.
+    assert not f().messages

mitmproxy/proxy/protocol2/test/tutils.py

@@ -3,8 +3,6 @@ import difflib
 import itertools
 import typing
 
-import re
-
 from mitmproxy.proxy.protocol2 import commands
 from mitmproxy.proxy.protocol2 import events
 from mitmproxy.proxy.protocol2 import layer
@@ -124,8 +122,9 @@ class playbook:
             self._errored = True
 
             def _str(x):
-                x_str = re.sub(r'Placeholder\((.*?)\)', r'\1', str(x))
+                arrow = ">" if isinstance(x, events.Event) else "<"
-                return f"{'>' if isinstance(x, events.Event) else '<'} {x_str}"
+                x = str(x).replace('Placeholder:', '')
+                return f"{arrow} {x}"
 
             diff = "\n".join(difflib.ndiff(
                 [_str(x) for x in self.expected],
@@ -136,15 +135,34 @@ class playbook:
             return True
 
     def __del__(self):
+        # Playbooks are only executed on assert (which signals that the playbook is partially
+        # complete), so we need to signal if someone forgets to assert and playbooks aren't
+        # evaluated.
         if not self._errored and len(self.actual) < len(self.expected):
             raise RuntimeError("Unfinished playbook!")
 
     def fork(self):
+        """
+        Fork the current playbook to assert a second execution stream from here on.
+        Returns a new playbook instance.
+        """
         return copy.deepcopy(self)
 
 
 class Placeholder:
-    """Placeholder value in playbooks, so that flows can be referenced before they are initialized."""
+    """
+    Placeholder value in playbooks, so that objects (flows in particular) can be referenced before
+    they are known. Example:
+
+    f = Placeholder()
+    assert (
+        playbook(tcp.TCPLayer(tctx))
+        << commands.Hook("tcp_start", f)  # the flow object returned here is generated by the layer.
+    )
+
+    # We can obtain the flow object now using f():
+    assert f().messages == 0
+    """
 
     def __init__(self):
         self.obj = None
@@ -154,4 +172,4 @@ class Placeholder:
         return self.obj
 
     def __repr__(self):
-        return f"Placeholder({repr(self.obj)})"
+        return f"Placeholder:{repr(self.obj)}"